Double Pendulum example (#17)

foldfelis · yuehhua · web-flow · commit f502087783c2 · 2021-08-28T01:52:24.000+08:00
* add DoublePendulumChaotic dataset

* build proj for double pendulum

* get data from file

* data visualization

* add annotation

* get_dataloader

* revise get_dataloader

* normalize data

* implement MNO

* revise model

* Update example/DoublePendulum/test/data.jl

Co-authored-by: Yueh-Hua Tu &lt;a504082002@gmail.com&gt;

* Update README.md

Co-authored-by: Yueh-Hua Tu &lt;a504082002@gmail.com&gt;

* Update README.md

Co-authored-by: Yueh-Hua Tu &lt;a504082002@gmail.com&gt;

* revise model

* revise model and save model

* fallback to last model

* train MNO in 2-D and make sure loss&lt;1e-2 about 5e-3 at epoch=35

* implement get_model and refactor

* result visualization (failed)

* revise data

* add gradient

* implement MNO

* fix test

* BPTT not work QQ

* navier stokes equation

* implement MNO on flow over a circle

* bug fix

* implement demo

* add MNO to model

* update readme

Co-authored-by: Yueh-Hua Tu &lt;a504082002@gmail.com&gt;
diff --git a/.gitignore b/.gitignore
@@ -22,3 +22,5 @@ docs/site/
 # committed for packages, but should be committed for applications that require a static
 # environment.
 Manifest.toml
+
+*.jld2
diff --git a/README.md b/README.md
@@ -20,7 +20,8 @@ between two continuous function spaces. The kernel can be trained on different g
 
 Fourier neural operator learns a neural operator with Dirichlet kernel to form a Fourier transformation. It performs Fourier transformation across infinite-dimensional function spaces and learns better than neural operator.
 
-Currently, `FourierOperator` is provided in this work.
+Currently, the `FourierOperator` layer is provided in this work.
+As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided. Please take a glance at them [here](src/model.jl).
 
 ## Usage
 
@@ -68,25 +69,34 @@ PDE training examples are provided in `example` folder.
 Use following commend to train model:
 
 ```julia
-$ julia --proj=example/Burgers
+$ julia --proj
 
 julia> using Burgers; Burgers.train()
 ```
 
-### Two-dimensional Darcy flow equation
+### Two-dimensional with time Navier-Stokes equation
 
-WIP
+The Navier-Stokes equation is learned by the `MarkovNeuralOperator` with only one time step information. Example can be found in `example/FlowOverCircle`.
 
-### Two-dimensional Navier-Stokes equation
+| **Ground Truth** | **Inferenced** |
+|:----------------:|:--------------:|
+| ![](example/FlowOverCircle/gallery/ans.gif) | ![](example/FlowOverCircle/gallery/inferenced.gif) |
 
-WIP
+Use following commend to train model:
+
+```julia
+$ julia --proj
+
+julia> using FlowOverCircle; FlowOverCircle.train()
+```
 
 ## Roadmap
 
 - [x] `FourierOperator` layer
 - [x] One-dimensional Burgers' equation example
-- [ ] Two-dimensional Darcy flow equation example
-- [ ] Two-dimensional Navier-Stokes equation example
+- [x] Two-dimensional with time Navier-Stokes equations example
+- [x] `MarkovNeuralOperator` model
+- [x] Flow over a circle prediction example
 - [ ] `NeuralOperator` layer
 - [ ] Poisson equation example
 
diff --git a/example/Burgers/src/data.jl b/example/Burgers/src/data.jl
@@ -2,8 +2,6 @@ using DataDeps
 using Fetch
 using MAT
 
-export get_burgers_data
-
 function register_burgers()
     register(DataDep(
         "Burgers",
@@ -18,7 +16,7 @@ function register_burgers()
     ))
 end
 
-function get_burgers_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples), T=Float32)
+function get_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples), T=Float32)
     file = matopen(joinpath(datadep"Burgers", "burgers_data_R10.mat"))
     x_data = T.(collect(read(file, "a")[1:n, 1:Δsamples:end]'))
     y_data = T.(collect(read(file, "u")[1:n, 1:Δsamples:end]'))
@@ -32,7 +30,7 @@ function get_burgers_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples
 end
 
 function get_dataloader(; n_train=1800, n_test=200, batchsize=100)
-    𝐱, 𝐲 = get_burgers_data(n=2048)
+    𝐱, 𝐲 = get_data(n=2048)
 
     𝐱_train, 𝐲_train = 𝐱[:, :, 1:n_train], 𝐲[:, 1:n_train]
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
diff --git a/example/Burgers/test/data.jl b/example/Burgers/test/data.jl
@@ -1,5 +1,5 @@
 @testset "get burgers data" begin
-    xs, ys = get_burgers_data(n=1000)
+    xs, ys = Burgers.get_data(n=1000)
 
     @test size(xs) == (2, 1024, 1000)
     @test size(ys) == (1024, 1000)
diff --git a/example/FlowOverCircle/Project.toml b/example/FlowOverCircle/Project.toml
@@ -0,0 +1,19 @@
+name = "FlowOverCircle"
+uuid = "1fc04e5d-1dd1-42ff-8d75-1d53504b2476"
+
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+WaterLily = "ed894a53-35f9-47f1-b17f-85db9237eebd"
+
+[extras]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test"]
diff --git a/example/FlowOverCircle/gallery/ans.gif b/example/FlowOverCircle/gallery/ans.gif
diff --git a/example/FlowOverCircle/gallery/inferenced.gif b/example/FlowOverCircle/gallery/inferenced.gif
diff --git a/example/FlowOverCircle/model/.gitkeep b/example/FlowOverCircle/model/.gitkeep
diff --git a/example/FlowOverCircle/notebook/mno.jl b/example/FlowOverCircle/notebook/mno.jl
@@ -0,0 +1,89 @@
+### A Pluto.jl notebook ###
+# v0.15.1
+
+using Markdown
+using InteractiveUtils
+
+# ╔═╡ 194baef2-0417-11ec-05ab-4527ef614024
+using Pkg; Pkg.develop(path=".."); Pkg.activate("..")
+
+# ╔═╡ 38c9ced5-dcf8-4e03-ac07-7c435687861b
+using FlowOverCircle, Plots
+
+# ╔═╡ 50ce80a3-a1e8-4ba9-a032-dad315bcb432
+md"
+# Markov Neural Operator
+
+JingYu Ning
+"
+
+# ╔═╡ 59769504-ebd5-4c6f-981f-d03826d8e34a
+md"
+This demo trains a Markov neural operator (MNO) introduced by [Zongyi Li *et al.*](https://arxiv.org/abs/2106.06898) with only one time step information. Then composed the operator to a Markov chain and inference the Navier-Stokes equations."
+
+# ╔═╡ 823b3547-6723-43cf-85e6-cc6eb44efea1
+md"
+## Generate data
+"
+
+# ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
+begin 
+	n = 20
+	data = FlowOverCircle.gen_data(LinRange(100, 100+n-1, n))
+end;
+
+# ╔═╡ 9b02b6a2-33c3-4ca6-bfba-0bd74b664830
+begin
+	anim = @animate for i in 1:size(data)[end]
+		heatmap(data[1, :, :, i]', color=:coolwarm, clim=(-1.5, 1.5))
+		scatter!(
+			[size(data, 3)÷2], [size(data, 3)÷2-1], 
+			markersize=45, color=:black, legend=false, ticks=false
+		)
+		annotate!(5, 5, text("i=$i", :left))
+	end
+	gif(anim, fps=2)
+end
+
+# ╔═╡ 55058635-c7e9-4ee3-81c2-0153e84f4c8e
+md"
+## Inference
+
+Use the first data generated above as the initial state, and apply the operator recurrently.
+"
+
+# ╔═╡ fbc287b8-f232-4350-9948-2091908e5a30
+begin
+	m = FlowOverCircle.get_model()
+	
+	states = Array{Float32}(undef, size(data))
+	states[:, :, :, 1] .= view(data, :, :, :, 1)
+	for i in 2:size(data)[end]
+		states[:, :, :, i:i] .= m(view(states, :, :, :, i-1:i-1))
+	end
+end
+
+# ╔═╡ a0b5e94c-a839-4cc0-a325-1a4ac39fafbc
+begin
+	anim_model = @animate for i in 1:size(states)[end]
+		heatmap(states[1, :, :, i]', color=:coolwarm, clim=(-1.5, 1.5))
+		scatter!(
+			[size(data, 3)÷2], [size(data, 3)÷2-1], 
+			markersize=45, color=:black, legend=false, ticks=false
+		)
+		annotate!(5, 5, text("i=$i", :left))
+	end
+	gif(anim_model, fps=2)
+end
+
+# ╔═╡ Cell order:
+# ╟─50ce80a3-a1e8-4ba9-a032-dad315bcb432
+# ╟─59769504-ebd5-4c6f-981f-d03826d8e34a
+# ╟─194baef2-0417-11ec-05ab-4527ef614024
+# ╠═38c9ced5-dcf8-4e03-ac07-7c435687861b
+# ╟─823b3547-6723-43cf-85e6-cc6eb44efea1
+# ╠═5268feee-bda2-4612-9d4c-a1db424a11c7
+# ╟─9b02b6a2-33c3-4ca6-bfba-0bd74b664830
+# ╟─55058635-c7e9-4ee3-81c2-0153e84f4c8e
+# ╠═fbc287b8-f232-4350-9948-2091908e5a30
+# ╟─a0b5e94c-a839-4cc0-a325-1a4ac39fafbc
diff --git a/example/FlowOverCircle/src/FlowOverCircle.jl b/example/FlowOverCircle/src/FlowOverCircle.jl
@@ -0,0 +1,63 @@
+module FlowOverCircle
+
+using NeuralOperators
+using Flux
+using CUDA
+using JLD2
+
+include("data.jl")
+
+function update_model!(model_file_path, model)
+    model = cpu(model)
+    jldsave(model_file_path; model)
+    @warn "model updated!"
+end
+
+function train()
+    if has_cuda()
+        @info "CUDA is on"
+        device = gpu
+        CUDA.allowscalar(false)
+    else
+        device = cpu
+    end
+
+    m = Chain(
+        Dense(1, 64),
+        FourierOperator(64=>64, (24, 24), gelu),
+        FourierOperator(64=>64, (24, 24), gelu),
+        FourierOperator(64=>64, (24, 24), gelu),
+        FourierOperator(64=>64, (24, 24), gelu),
+        Dense(64, 1),
+    ) |> device
+
+    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+
+    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+
+    @info "gen data... "
+    @time loader_train, loader_test = get_dataloader()
+
+    losses = Float32[]
+    function validate()
+        validation_loss = sum(loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test)/length(loader_test)
+        @info "loss: $validation_loss"
+
+        push!(losses, validation_loss)
+        (losses[end] == minimum(losses)) && update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
+    end
+    call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
+
+    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
+    Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
+end
+
+function get_model()
+    f = jldopen(joinpath(@__DIR__, "../model/model.jld2"))
+    model = f["model"]
+    close(f)
+
+    return model
+end
+
+end
diff --git a/example/FlowOverCircle/src/data.jl b/example/FlowOverCircle/src/data.jl
@@ -0,0 +1,43 @@
+using WaterLily
+using LinearAlgebra: norm2
+
+"""
+    circle(n, m; Re=250)
+
+This function is copy from [WaterLily](https://github.com/weymouth/WaterLily.jl)
+"""
+function circle(n, m; Re=250)
+    # Set physical parameters
+    U, R, center = 1., m/8., [m/2, m/2]
+    ν = U * R / Re
+
+    body = AutoBody((x,t) -> norm2(x .- center) - R)
+    Simulation((n+2, m+2), [U, 0.], R; ν, body)
+end
+
+function gen_data(ts::AbstractRange)
+    n, m = 3(2^5), 2^6
+    circ = circle(n, m)
+
+    𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
+    for (i, t) in enumerate(ts)
+        sim_step!(circ, t)
+        𝐩s[:, :, :, i] = Float32.(circ.flow.p)[2:end-1, 2:end-1]
+    end
+
+    return 𝐩s
+end
+
+function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Float64=0.95, batchsize=100)
+    data = gen_data(ts)
+
+    n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))
+
+    𝐱_train, 𝐲_train = data[:, :, :, 1:(n_train-1)], data[:, :, :, 2:n_train]
+    loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
+
+    𝐱_test, 𝐲_test = data[:, :, :, (end-n_test+1):(end-1)], data[:, :, :, (end-n_test+2):end]
+    loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
+
+    return loader_train, loader_test
+end
diff --git a/example/FlowOverCircle/test/data.jl b/example/FlowOverCircle/test/data.jl
@@ -0,0 +1,3 @@
+@testset "flow over circle" begin
+
+end
diff --git a/example/FlowOverCircle/test/runtests.jl b/example/FlowOverCircle/test/runtests.jl
@@ -0,0 +1,6 @@
+using FlowOverCircle
+using Test
+
+@testset "FlowOverCircle" begin
+    include("data.jl")
+end
diff --git a/src/model.jl b/src/model.jl
@@ -1,5 +1,6 @@
 export
-    FourierNeuralOperator
+    FourierNeuralOperator,
+    MarkovNeuralOperator
 
 function FourierNeuralOperator(;
     ch=(2, 64, 64, 64, 64, 64, 128, 1),
@@ -17,3 +18,18 @@ function FourierNeuralOperator(;
         flatten
     )
 end
+
+function MarkovNeuralOperator(;
+    ch=(1, 64, 64, 64, 64, 64, 1),
+    modes=(24, 24),
+    σ=gelu
+)
+    return Chain(
+        Dense(ch[1], ch[2]),
+        FourierOperator(ch[2]=>ch[3], modes, σ),
+        FourierOperator(ch[3]=>ch[4], modes, σ),
+        FourierOperator(ch[4]=>ch[5], modes, σ),
+        FourierOperator(ch[5]=>ch[6], modes, σ),
+        Dense(ch[6], ch[7])
+    )
+end
diff --git a/test/model.jl b/test/model.jl
@@ -5,6 +5,17 @@
     @test size(m(𝐱)) == size(𝐲)
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
-    data = [(𝐱[:, :, 1:5], 𝐲[:, 1:5])]
+    data = [(𝐱, 𝐲)]
+    Flux.train!(loss, params(m), data, Flux.ADAM())
+end
+
+@testset "MarkovNeuralOperator" begin
+    m = MarkovNeuralOperator()
+
+    𝐱, 𝐲 = rand(Float32, 1, 64, 64, 5), rand(Float32, 1, 64, 64, 5)
+    @test size(𝐱) == size(m(𝐱)) == size(𝐲)
+
+    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+    data = [(𝐱, 𝐲)]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
