Skip to content

Prevent inconsistent indicators when growing batched graphs #737

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DHAKED440 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Prevent inconsistent indicators when growing batched graphs #737

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion GNNGraphs/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "GNNGraphs"
uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
version = "1.6.1-DEV"
version = "1.6.2-DEV"
authors = ["Carlo Lucibello and contributors"]

[workspace]
Expand Down
24 changes: 19 additions & 5 deletions GNNGraphs/src/transform.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,15 @@ function remove_nodes(g::GNNGraph, p::AbstractFloat)
return remove_nodes(g, nodes_to_remove)
end

function _graph_indicator_after_node_growth(g::GNNGraph, n_new_nodes::Integer)
n_new_nodes <= 0 && return g.graph_indicator
g.num_graphs > 1 && throw(ArgumentError(
"Cannot add nodes to a batched graph without specifying which graph they belong to.",
))
isnothing(g.graph_indicator) && return nothing
return cat_features(g.graph_indicator, ones_like(g.graph_indicator, n_new_nodes))
end

"""
add_edges(g::GNNGraph, s::AbstractVector, t::AbstractVector; [edata])
add_edges(g::GNNGraph, (s, t); [edata])
Expand All @@ -290,7 +299,8 @@ Optionally, pass the edge weight `w` and the features `edata` for the new edges
Returns a new graph sharing part of the underlying data with `g`.

If the `s` or `t` contain nodes that are not already present in the graph,
they are added to the graph as well.
they are added to the graph as well. Adding new nodes to a batched graph is
not supported because their graph assignment is ambiguous.

# Examples

Expand Down Expand Up @@ -333,6 +343,9 @@ function add_edges(g::GNNGraph{<:COO_T}, data::COO_T; edata = nothing)
end
@assert minimum(snew) >= 1
@assert minimum(tnew) >= 1
num_nodes = max(maximum(snew), maximum(tnew), g.num_nodes)
graph_indicator = _graph_indicator_after_node_growth(g, num_nodes - g.num_nodes)

num_new = length(snew)
edata = normalize_graphdata(edata, default_name = :e, n = num_new)
edata = cat_features(g.edata, edata)
Expand All @@ -343,7 +356,6 @@ function add_edges(g::GNNGraph{<:COO_T}, data::COO_T; edata = nothing)
w = get_edge_weight(g)
w = cat_features(w, wnew, g.num_edges, num_new)

num_nodes = max(maximum(snew), maximum(tnew), g.num_nodes)
if num_nodes > g.num_nodes
ndata_new = normalize_graphdata((;), default_name = :x, n = num_nodes - g.num_nodes)
ndata = cat_features(g.ndata, ndata_new)
Expand All @@ -353,7 +365,7 @@ function add_edges(g::GNNGraph{<:COO_T}, data::COO_T; edata = nothing)

return GNNGraph((s, t, w),
num_nodes, length(s), g.num_graphs,
g.graph_indicator,
graph_indicator,
ndata, edata, g.gdata)
end

Expand Down Expand Up @@ -553,15 +565,17 @@ end

Add `n` new nodes to graph `g`. In the
new graph, these nodes will have indexes from `g.num_nodes + 1`
to `g.num_nodes + n`.
to `g.num_nodes + n`. Adding nodes to a batched graph is not supported
because their graph assignment is ambiguous.
"""
function add_nodes(g::GNNGraph{<:COO_T}, n::Integer; ndata = (;))
graph_indicator = _graph_indicator_after_node_growth(g, n)
ndata = normalize_graphdata(ndata, default_name = :x, n = n)
ndata = cat_features(g.ndata, ndata)

GNNGraph(g.graph,
g.num_nodes + n, g.num_edges, g.num_graphs,
g.graph_indicator,
graph_indicator,
ndata, g.edata, g.gdata)
end

Expand Down
34 changes: 34 additions & 0 deletions GNNGraphs/test/transform.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,6 +467,40 @@ end
end
end

@testitem "node growth and batched graph indicators" setup=[GraphsTestModule] begin
using .GraphsTestModule

g1 = GNNGraph(([1, 2], [2, 3]); num_nodes = 3)
g2 = GNNGraph(([1, 2], [2, 1]); num_nodes = 4)
g = batch([g1, g2])

@test g.num_nodes == length(g.graph_indicator) == 7
@test_throws ArgumentError add_nodes(g, 2)
@test_throws ArgumentError add_edges(g, [8], [1])
@test add_nodes(g, 0).graph_indicator == g.graph_indicator
@test add_edges(g, Int[], Int[]) === g

# Adding edges between existing nodes does not grow the batched graph.
g_with_edge = add_edges(g, [1], [2])
@test g_with_edge.num_nodes == length(g_with_edge.graph_indicator) == 7

# A single graph with explicit membership can safely extend that membership.
single = GNNGraph(([1], [2]); num_nodes = 2, graph_indicator = [1, 1])
single_with_nodes = add_nodes(single, 2)
single_with_edge = add_edges(single, [3], [1])
@test single_with_nodes.graph_indicator == [1, 1, 1, 1]
@test single_with_edge.graph_indicator == [1, 1, 1]

# Unbatched graphs represent membership implicitly, so their stored field stays `nothing`.
unbatched = GNNGraph(([1], [2]); num_nodes = 2)
unbatched_with_nodes = add_nodes(unbatched, 2)
unbatched_with_edge = add_edges(unbatched, [3], [1])
@test isnothing(unbatched_with_nodes.graph_indicator)
@test length(graph_indicator(unbatched_with_nodes)) == unbatched_with_nodes.num_nodes
@test isnothing(unbatched_with_edge.graph_indicator)
@test length(graph_indicator(unbatched_with_edge)) == unbatched_with_edge.num_nodes
end

@testitem "remove_self_loops" setup=[GraphsTestModule] begin
using .GraphsTestModule
for GRAPH_T in GRAPH_TYPES
Expand Down
Loading