Skip to content
This repository was archived by the owner on Sep 28, 2024. It is now read-only.

Commit 919554e

Browse files
authored
Merge pull request #54 from yuehhua/gno
Implement Graph Neural Operator (GNO)
2 parents 8ca0713 + e9b107e commit 919554e

File tree

13 files changed

+291
-83
lines changed

13 files changed

+291
-83
lines changed

Project.toml

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ CUDAKernels = "72cfdca4-0801-4ab0-bf6a-d52aa10adc57"
99
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
1010
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
1111
Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
12+
GeometricFlux = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
1213
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
14+
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
1315
Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
1416
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
1517

@@ -19,13 +21,15 @@ CUDAKernels = "0.3, 0.4"
1921
ChainRulesCore = "1.13"
2022
FFTW = "1.4"
2123
Flux = "0.12"
24+
GeometricFlux = "0.10"
2225
KernelAbstractions = "0.7, 0.8"
2326
Tullio = "0.3"
2427
Zygote = "0.6"
2528
julia = "1.6"
2629

2730
[extras]
31+
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
2832
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
2933

3034
[targets]
31-
test = ["Test"]
35+
test = ["Graphs", "Test"]

docs/make.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ makedocs(;
1515
),
1616
pages=[
1717
"Home" => "index.md",
18+
"Introduction" => "introduction.md",
1819
"APIs" => "apis.md"
1920
],
2021
)

docs/src/apis.md

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
1-
## Index
2-
3-
```@index
4-
```
1+
# APIs
52

63
## Layers
74

@@ -40,6 +37,25 @@ OperatorKernel
4037

4138
Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
4239

40+
---
41+
42+
### Graph kernel layer
43+
44+
```math
45+
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) \} )
46+
```
47+
48+
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``.
49+
Activation function ``\sigma`` can be arbitrary non-linear function.
50+
51+
```@docs
52+
GraphKernel
53+
```
54+
55+
Reference: [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://arxiv.org/abs/2003.03485)
56+
57+
---
58+
4359
## Models
4460

4561
### Fourier neural operator

docs/src/index.md

Lines changed: 4 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,40 +4,13 @@ CurrentModule = NeuralOperators
44

55
# NeuralOperators
66

7-
Documentation for [NeuralOperators](https://github.com/foldfelis/NeuralOperators.jl).
8-
9-
| **Ground Truth** | **Inferenced** |
10-
|:----------------:|:--------------:|
117
| ![](https://github.com/foldfelis/NeuralOperators.jl/blob/master/example/FlowOverCircle/gallery/ans.gif?raw=true) | ![](https://github.com/foldfelis/NeuralOperators.jl/blob/master/example/FlowOverCircle/gallery/inferenced.gif?raw=true) |
8+
|:----------------:|:--------------:|
9+
| **Ground Truth** | **Inferenced** |
1210

1311
The demonstration shown above is Navier-Stokes equation learned by the `MarkovNeuralOperator` with only one time step information.
1412
Example can be found in [`example/FlowOverCircle`](https://github.com/foldfelis/NeuralOperators.jl/tree/master/example/FlowOverCircle).
1513

16-
## Abstract
17-
18-
Neural operator is a novel deep learning architecture.
19-
It learns a operator, which is a mapping between infinite-dimensional function spaces.
20-
It can be used to resolve [partial differential equations (PDE)](https://en.wikipedia.org/wiki/Partial_differential_equation).
21-
Instead of solving by finite element method, a PDE problem can be resolved by training a neural network to learn an operator mapping
22-
from infinite-dimensional space (u, t) to infinite-dimensional space f(u, t).
23-
Neural operator learns a continuous function between two continuous function spaces.
24-
The kernel can be trained on different geometry, which is learned from a graph.
25-
26-
**Fourier neural operator** learns a neural operator with Dirichlet kernel to form a Fourier transformation.
27-
It performs Fourier transformation across infinite-dimensional function spaces and learns better than neural operator.
28-
29-
**Markov neural operator** learns a neural operator with Fourier operators.
30-
With only one time step information of learning, it can predict the following few steps with low loss
31-
by linking the operators into a Markov chain.
32-
33-
**DeepONet operator** (Deep Operator Network) learns a neural operator with the help of two sub-neural net structures described as the branch and the trunk network.
34-
The branch network is fed the initial conditions data, whereas the trunk is fed with the locations where the target(output) is evaluated from the corresponding initial conditions.
35-
It is important that the output size of the branch and trunk subnets is same so that a dot product can be performed between them.
36-
37-
Currently, the `OperatorKernel` layer is provided in this work.
38-
As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided.
39-
Please take a glance at them [here](apis.html#Models).
40-
4114
## Quick start
4215

4316
The package can be installed with the Julia package manager. From the Julia REPL, type `]` to enter the Pkg REPL mode and run:
@@ -80,7 +53,7 @@ model = FourierNeuralOperator(
8053
And then train as a Flux model.
8154

8255
```julia
83-
loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- model(𝐱)) / size(𝐱)[end]
56+
loss(𝐱, 𝐲) = Flux.Losses.mse(model(𝐱), 𝐲)
8457
opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
8558
Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)
8659
```
@@ -112,4 +85,4 @@ opt = ADAM(learning_rate)
11285
parameters = params(model)
11386
Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb=evalcb)
11487
```
115-
A more complete example using DeepONet architecture to solve Burgers' equation can be found in the [examples](../../example/Burgers/src/Burgers_deeponet.jl)
88+
A more complete example using DeepONet architecture to solve Burgers' equation can be found in the [examples](https://github.com/foldfelis/NeuralOperators.jl/tree/master/example/Burgers/src/Burgers_deeponet.jl).

docs/src/introduction.md

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Introduction
2+
3+
Neural operator is a novel deep learning architecture.
4+
It learns a operator, which is a mapping between infinite-dimensional function spaces.
5+
It can be used to resolve [partial differential equations (PDE)](https://en.wikipedia.org/wiki/Partial_differential_equation).
6+
Instead of solving by time-consuming finite element method, a PDE problem can be resolved by training a neural network to learn
7+
an operator mapping from infinite-dimensional space ``(u, t)`` to infinite-dimensional space ``f(u, t)``.
8+
Neural operator learns a continuous function between two continuous function spaces.
9+
The kernel can be trained on different geometry, including regular Euclidean space or a graph topology.
10+
11+
## Fourier Neural Operators
12+
13+
Fourier neural operator (FNO) learns a neural operator with Dirichlet kernel to form a Fourier transformation.
14+
It performs Fourier transformation across infinite-dimensional function spaces and learns better than neural operator.
15+
16+
## Markov Neural Operators
17+
18+
Markov neural operator (MNO) learns a neural operator with Fourier operators.
19+
With only one time step information of learning, it can predict the following few steps with low loss
20+
by linking the operators into a Markov chain.
21+
22+
## Deep Operator Network
23+
24+
Deep operator network (DeepONet) learns a neural operator with the help of two sub-neural network structures described as the branch and the trunk network.
25+
The branch network is fed the initial conditions data, whereas the trunk is fed with the locations where the target(output) is evaluated from the corresponding initial conditions.
26+
It is important that the output size of the branch and trunk subnets is same so that a dot product can be performed between them.
27+
28+
Currently, the `OperatorKernel` layer is provided in this work.
29+
As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided.
30+
Please take a glance at [them](apis.html#Models).

example/SuperResolution/Project.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@ uuid = "a8258e1f-331c-4af2-83e9-878628278453"
44
[deps]
55
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
66
Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
7+
GeometricFlux = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
8+
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
79
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
810
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
911
NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
1012
Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
1113
Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
14+
ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
1215
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
1316
WaterLily = "ed894a53-35f9-47f1-b17f-85db9237eebd"
1417

example/SuperResolution/src/SuperResolution.jl

Lines changed: 10 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -2,54 +2,21 @@ module SuperResolution
22

33
using NeuralOperators
44
using Flux
5+
using Flux.Losses: mse
6+
using Flux.Data: DataLoader
7+
using GeometricFlux
8+
using Graphs
59
using CUDA
610
using JLD2
11+
using ProgressMeter: Progress, next!
712

813
include("data.jl")
14+
include("models.jl")
915

1016
function update_model!(model_file_path, model)
1117
model = cpu(model)
1218
jldsave(model_file_path; model)
13-
@warn "model updated!"
14-
end
15-
16-
function train()
17-
if has_cuda()
18-
@info "CUDA is on"
19-
device = gpu
20-
CUDA.allowscalar(false)
21-
else
22-
device = cpu
23-
end
24-
25-
m = Chain(
26-
Dense(1, 64),
27-
OperatorKernel(64=>64, (24, 24), gelu),
28-
OperatorKernel(64=>64, (24, 24), gelu),
29-
OperatorKernel(64=>64, (24, 24), gelu),
30-
OperatorKernel(64=>64, (24, 24), gelu),
31-
Dense(64, 1),
32-
) |> device
33-
34-
loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
35-
36-
opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
37-
38-
@info "gen data... "
39-
@time loader_train, loader_test = get_dataloader()
40-
41-
losses = Float32[]
42-
function validate()
43-
validation_loss = sum(loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test)/length(loader_test)
44-
@info "loss: $validation_loss"
45-
46-
push!(losses, validation_loss)
47-
(losses[end] == minimum(losses)) && update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
48-
end
49-
call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
50-
51-
data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
52-
Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
19+
@info "model updated!"
5320
end
5421

5522
function get_model()
@@ -60,4 +27,7 @@ function get_model()
6027
return model
6128
end
6229

30+
loss(m, 𝐱, 𝐲) = mse(m(𝐱), 𝐲)
31+
loss(m, loader::DataLoader, device) = sum(loss(m, 𝐱 |> device, 𝐲 |> device) for (𝐱, 𝐲) in loader)/length(loader)
32+
6333
end

example/SuperResolution/src/data.jl

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,20 @@ function circle(n, m; Re=250)
1515
Simulation((n+2, m+2), [U, 0.], R; ν, body)
1616
end
1717

18-
function gen_data(ts::AbstractRange)
18+
function gen_data(ts::AbstractRange, T=Float32)
1919
n, m = 2 * 3(2^5), 2 * 2^6
2020
circ = circle(n, m)
2121

22-
𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
22+
𝐩s = Array{T}(undef, 1, n, m, length(ts))
2323
for (i, t) in enumerate(ts)
2424
sim_step!(circ, t)
25-
𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
25+
𝐩s[1, :, :, i] .= T.(circ.flow.p)[2:end-1, 2:end-1]
2626
end
2727

2828
return 𝐩s
2929
end
3030

31-
function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Float64=0.95, batchsize=100)
31+
function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Real=0.95, batchsize=100)
3232
data = gen_data(ts)
3333

3434
n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))
@@ -41,3 +41,19 @@ function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::
4141

4242
return loader_train, loader_test
4343
end
44+
45+
function get_same_resolution(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Real=0.95, batchsize=100)
46+
data = gen_data(ts)
47+
48+
n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))
49+
50+
𝐱_train, 𝐲_train = data[:, 1:2:end, 1:2:end, 1:(n_train-1)], data[:, 1:2:end, 1:2:end, 2:n_train]
51+
𝐱_train, 𝐲_train = reshape(𝐱_train, 1, :, n_train-1), reshape(𝐲_train, 1, :, n_train-1)
52+
loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
53+
54+
𝐱_test, 𝐲_test = data[:, 1:2:end, 1:2:end, (end-n_test+1):(end-1)], data[:, 1:2:end, 1:2:end, (end-n_test+2):end]
55+
𝐱_test, 𝐲_test = reshape(𝐱_test, 1, :, n_test-1), reshape(𝐲_test, 1, :, n_test-1)
56+
loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
57+
58+
return loader_train, loader_test
59+
end

0 commit comments

Comments
 (0)