add GraphKernel

add doc

fix

Author: yuehhua

Project.toml

@@ -9,7 +9,9 @@ CUDAKernels = "72cfdca4-0801-4ab0-bf6a-d52aa10adc57"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+GeometricFlux = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -19,13 +21,15 @@ CUDAKernels = "0.3, 0.4"
 ChainRulesCore = "1.13"
 FFTW = "1.4"
 Flux = "0.12"
+GeometricFlux = "0.10"
 KernelAbstractions = "0.7, 0.8"
 Tullio = "0.3"
 Zygote = "0.6"
 julia = "1.6"
 
 [extras]
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Graphs", "Test"]

docs/src/apis.md

@@ -40,6 +40,21 @@ OperatorKernel
 
 Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
 
+### Graph kernel layer
+
+```math
+v_{t+1}(x_i) = \sigma(W v_t(x_i) + \frac{1}{|\mathcal{N}(x_i)|} \sum_{x_j \in \mathcal{N}(x_i)} \kappa \{ v_t(x_i), v_t(x_j) \} )
+```
+
+where ``v_t(x_i)`` is the input function for ``t``-th layer, ``x_i`` is the node feature for ``i``-th node and ``\mathcal{N}(x_i)`` represents the neighbors for ``x_i``.
+Activation function ``\sigma`` can be arbitrary non-linear function.
+
+```@docs
+GraphKernel
+```
+
+Reference: [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://arxiv.org/abs/2003.03485)
+
 ## Models
 
 ### Fourier neural operator

src/NeuralOperators.jl

@@ -7,6 +7,8 @@ module NeuralOperators
     using KernelAbstractions
     using Zygote
     using ChainRulesCore
+    using GeometricFlux
+    using Statistics
 
     export DeepONet
 

src/operator_kernel.jl

@@ -1,7 +1,8 @@
 export
     OperatorConv,
     SpectralConv,
-    OperatorKernel
+    OperatorKernel,
+    GraphKernel
 
 struct OperatorConv{P, T, S, TT}
     weight::T
@@ -170,6 +171,52 @@ function (m::OperatorKernel)(𝐱)
     return m.σ.(m.linear(𝐱) + m.conv(𝐱))
 end
 
+"""
+    GraphKernel(κ, ch, σ=identity)
+
+Graph kernel layer.
+
+## Arguments
+
+* `κ`: A neural network layer for approximation, e.g. a `Dense` layer or a MLP.
+* `ch`: Channel size for linear transform, e.g. `32`.
+* `σ`: Activation function.
+"""
+struct GraphKernel{A,B,F} <: MessagePassing
+    linear::A
+    κ::B
+    σ::F
+end
+
+function GraphKernel(κ, ch::Int, σ=identity; init=Flux.glorot_uniform)
+    W = init(ch, ch)
+    return GraphKernel(W, κ, σ)
+end
+
+Flux.@functor GraphKernel
+
+function GeometricFlux.message(l::GraphKernel, x_i::AbstractArray, x_j::AbstractArray, e_ij)
+    return l.κ(vcat(x_i, x_j))
+end
+
+function GeometricFlux.update(l::GraphKernel, m::AbstractArray, x::AbstractArray)
+    return l.σ.(GeometricFlux._matmul(l.linear, x) + m)
+end
+
+function (l::GraphKernel)(el::NamedTuple, X::AbstractArray)
+    GraphSignals.check_num_nodes(el.N, X)
+    _, V, _ = GeometricFlux.propagate(l, el, nothing, X, nothing, mean, nothing, nothing)
+    return V
+end
+
+function Base.show(io::IO, l::GraphKernel)
+    channel, _ = size(l.linear)
+    print(io, "GraphKernel(", l.κ, ", channel=", channel)
+    l.σ == identity || print(io, ", ", l.σ)
+    print(io, ")")
+end
+
+
 #########
 # utils #
 #########

test/operator_kernel.jl

@@ -154,3 +154,20 @@ end
     @test SpectralConv(ch, modes) isa OperatorConv
     @test SpectralConv(ch, modes).transform isa FourierTransform
 end
+
+@testset "GraphKernel" begin
+    batch_size = 5
+    channel = 32
+    N = 10*10
+
+    κ = Dense(2*channel, channel, relu)
+
+    graph = grid([10, 10])
+    𝐱 = rand(Float32, channel, N, batch_size)
+    l = WithGraph(FeaturedGraph(graph), GraphKernel(κ, channel))
+    @test repr(l.layer) == "GraphKernel(Dense(64, 32, relu), channel=32)"
+    @test size(l(𝐱)) == (channel, N, batch_size)
+
+    g = Zygote.gradient(() -> sum(l(𝐱)), Flux.params(l))
+    @test length(g.grads) == 3
+end

test/runtests.jl

@@ -1,7 +1,10 @@
 using NeuralOperators
-using Test
-using Flux
 using CUDA
+using Flux
+using GeometricFlux
+using Graphs
+using Zygote
+using Test
 
 CUDA.allowscalar(false)