Merge pull request #79 from yuehhua/gno

Fix GNO example

Author: yuehhua

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -3,6 +3,7 @@ module FlowOverCircle
 using WaterLily, LinearAlgebra, ProgressMeter, MLUtils
 using NeuralOperators, Flux, GeometricFlux, Graphs
 using CUDA, FluxTraining, BSON
+using GeometricFlux.GraphSignals: generate_grid
 
 function circle(n, m; Re = 250) # copy from [WaterLily](https://github.com/weymouth/WaterLily.jl)
     # Set physical parameters
@@ -31,16 +32,12 @@ function gen_data(ts::AbstractRange)
     return 𝐩s
 end
 
-function get_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
-                        ratio::Float64 = 0.95, batchsize = 100, flatten = false)
+function get_mno_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
+                            ratio::Float64 = 0.95, batchsize = 100)
     data = gen_data(ts)
     𝐱, 𝐲 = data[:, :, :, 1:(end - 1)], data[:, :, :, 2:end]
     n = length(ts) - 1
 
-    if flatten
-        𝐱, 𝐲 = reshape(𝐱, 1, :, n), reshape(𝐲, 1, :, n)
-    end
-
     data_train, data_test = splitobs(shuffleobs((𝐱, 𝐲)), at = ratio)
 
     loader_train = DataLoader(data_train, batchsize = batchsize, shuffle = true)
@@ -49,7 +46,7 @@ function get_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
     return loader_train, loader_test
 end
 
-function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
+function train_mno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
@@ -61,7 +58,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
 
     model = MarkovNeuralOperator(ch = (1, 64, 64, 64, 64, 64, 1), modes = (24, 24),
                                  σ = gelu)
-    data = get_dataloader()
+    data = get_mno_dataloader()
     optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
@@ -74,6 +71,32 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     return learner
 end
 
+function get_gno_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
+                            ratio::Float64 = 0.95, batchsize = 8)
+    data = gen_data(ts)
+    𝐱, 𝐲 = data[:, :, :, 1:(end - 1)], data[:, :, :, 2:end]
+    n = length(ts) - 1
+
+    # generate graph
+    graph = Graphs.grid(size(data)[2:3])
+
+    # add grid coordinates
+    grid = generate_coordinates(𝐱[1, :, :, 1])
+    grid = repeat(grid, outer = (1, 1, 1, n))
+    𝐱 = vcat(𝐱, grid)
+
+    # flatten
+    𝐱, 𝐲 = reshape(𝐱, size(𝐱, 1), :, n), reshape(𝐲, 1, :, n)
+
+    fg = FeaturedGraph(graph, nf = 𝐱, pf = 𝐱)
+    data_train, data_test = splitobs(shuffleobs((fg, 𝐲)), at = ratio)
+
+    loader_train = DataLoader(data_train, batchsize = batchsize, shuffle = true)
+    loader_test = DataLoader(data_test, batchsize = batchsize, shuffle = false)
+
+    return loader_train, loader_test
+end
+
 function train_gno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     if cuda && CUDA.has_cuda()
         device = gpu
@@ -84,17 +107,19 @@ function train_gno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
         @info "Training on CPU"
     end
 
-    featured_graph = FeaturedGraph(grid([96, 64]))
-    model = Chain(Dense(1, 16),
-                  WithGraph(featured_graph, GraphKernel(Dense(2 * 16, 16, gelu), 16)),
-                  WithGraph(featured_graph, GraphKernel(Dense(2 * 16, 16, gelu), 16)),
-                  WithGraph(featured_graph, GraphKernel(Dense(2 * 16, 16, gelu), 16)),
-                  WithGraph(featured_graph, GraphKernel(Dense(2 * 16, 16, gelu), 16)),
+    grid_dim = 2
+    edge_dim = 2(grid_dim + 1)
+    model = Chain(GraphParallel(node_layer = Dense(grid_dim + 1, 16)),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  node_feature,
                   Dense(16, 1))
-    data = get_dataloader(batchsize = 16, flatten = true)
+
     optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
-
+    data = get_gno_dataloader()
     learner = Learner(model, data, optimiser, loss_func,
                       ToDevice(device, device),
                       Checkpointer(joinpath(@__DIR__, "../model/")))

src/graph_kernel.jl

@@ -23,6 +23,7 @@ end
 
 function GraphKernel(κ, ch::Int, σ = identity; init = Flux.glorot_uniform)
     W = init(ch, ch)
+
     return GraphKernel(W, κ, σ)
 end
 
@@ -33,6 +34,7 @@ function GeometricFlux.message(l::GraphKernel, x_i, x_j::AbstractArray, e_ij::Ab
     K = l.κ(e_ij)
     dims = size(K)[2:end]
     m_ij = GeometricFlux._matmul(reshape(K, N, N, :), reshape(x_j, N, 1, :))
+
     return reshape(m_ij, N, dims...)
 end
 
@@ -42,13 +44,29 @@ end
 
 function (l::GraphKernel)(el::NamedTuple, X::AbstractArray, E::AbstractArray)
     GraphSignals.check_num_nodes(el.N, X)
-    GraphSignals.check_num_nodes(el.E, E)
+    GraphSignals.check_num_edges(el.E, E)
     _, V, _ = GeometricFlux.propagate(l, el, E, X, nothing, mean, nothing, nothing)
+
     return V
 end
 
+# For variable graph
+function (l::GraphKernel)(fg::AbstractFeaturedGraph)
+    nf = node_feature(fg)
+    pf = positional_feature(fg)
+    GraphSignals.check_num_nodes(fg, nf)
+    GraphSignals.check_num_nodes(fg, pf)
+    el = GeometricFlux.GraphSignals.to_namedtuple(fg)
+
+    # node features + positional features as edge features
+    ef = vcat(GeometricFlux._gather(pf, el.xs), GeometricFlux._gather(pf, el.nbrs))
+    _, V, _ = GeometricFlux.propagate(l, el, ef, nf, nothing, mean, nothing, nothing)
+
+    return ConcreteFeaturedGraph(fg, nf = V)
+end
+
 function Base.show(io::IO, l::GraphKernel)
-    channel, _ = size(l.linear)
+    channel = size(l.linear, 1)
     print(io, "GraphKernel(", l.κ, ", channel=", channel)
     l.σ == identity || print(io, ", ", l.σ)
     print(io, ")")

test/graph_kernel.jl

@@ -1,19 +1,24 @@
 @testset "GraphKernel" begin
     batch_size = 5
-    channel = 32
+    channel = 1
     coord_dim = 2
     N = 10
 
     graph = grid([N, N])
-    κ = Dense(2(coord_dim + 1), abs2(channel), relu)
-
     𝐱 = rand(Float32, channel, nv(graph), batch_size)
-    E = rand(Float32, 2(coord_dim + 1), ne(graph), batch_size)
-    l = WithGraph(FeaturedGraph(graph), GraphKernel(κ, channel))
-    @test repr(l.layer) ==
-          "GraphKernel(Dense($(2(coord_dim + 1)) => $(abs2(channel)), relu), channel=32)"
-    @test size(l(𝐱, E)) == (channel, nv(graph), batch_size)
+    κ = Dense(2(coord_dim + channel), abs2(channel), relu)
+    κ_in_dim, κ_out_dim = 2(coord_dim + channel), abs2(channel)
+
+    @testset "layer without graph" begin
+        pf = rand(Float32, coord_dim, nv(graph), batch_size)
+        pf = vcat(𝐱, pf)
+        l = GraphKernel(κ, channel)
+        fg = FeaturedGraph(graph, nf = 𝐱, pf = pf)
+        fg_ = l(fg)
+        @test size(node_feature(fg_)) == (channel, nv(graph), batch_size)
+        @test_throws MethodError l(𝐱)
 
-    g = Zygote.gradient(() -> sum(l(𝐱, E)), Flux.params(l))
-    @test length(g.grads) == 3
+        g = gradient(() -> sum(node_feature(l(fg))), Flux.params(l))
+        @test length(g.grads) == 5
+    end
 end