Merge pull request #89 from KirillZubov/burger_visualization

update Burger example

Author: yuehhua

example/Burgers/README.md

@@ -2,6 +2,15 @@
 
 In this example, a [Burgers' equation](https://en.wikipedia.org/wiki/Burgers%27_equation)
 is learned by a one-dimensional Fourier neural operator network.
+
+there is learn the operator mapping the initial condition to last point of time evolition of equation in some function space : 
+
+```math
+u(x, 0) ->  u(x, t_end)\
+```
+
+![](gallery/burgers.png) 
+
 Change directory to `example/Burgers` and use following commend to train model:
 
 ```julia

example/Burgers/gallery/burgers.png

@@ -0,0 +1,25 @@
+using Burgers, Plots
+using DataDeps, MAT, MLUtils
+using NeuralOperators, Flux
+using CUDA, BSON
+
+dataset = Burgers.get_data(n = 1000);
+
+m = Burgers.get_model();
+input_data, ground_truth = dataset[1], dataset[2];
+
+
+i = 1
+plot(input_data[1,:,1],ground_truth[1, :, i], label = "ground_truth",title = "                                              Burgers equation u(x,T_end)");
+p1 = plot!(input_data[1,:,1],m(view(input_data, :, :, i:i))[1, :, 1],label = "predict");
+plot(input_data[1,:,1],ground_truth[1, :, i + 1],label = "ground_truth");
+p2 = plot!(input_data[1,:,1],m(view(input_data, :, :, (i + 1):(i + 1)))[1, :, 1],label = "predict");
+i = 3
+
+plot(input_data[1,:,1],ground_truth[1, :, i], label = "ground_truth");
+p3 = plot!(input_data[1,:,1],m(view(input_data, :, :, i:i))[1, :, 1],label = "predict");
+plot(input_data[1,:,1],ground_truth[1, :, i + 1], label = "ground_truth");
+p4 = plot!(input_data[1,:,1],m(view(input_data, :, :, (i + 1):(i + 1)))[1, :, 1], label = "predict");
+p = plot(p1, p2, p3, p4)
+
+

example/Burgers/notebook/mno.jl

@@ -4,6 +4,7 @@ using DataDeps, MAT, MLUtils
 using NeuralOperators, Flux
 using CUDA, FluxTraining, BSON
 import Flux: params
+using BSON: @save, @load
 
 include("Burgers_deeponet.jl")
 
@@ -12,7 +13,14 @@ function register_burgers()
                      """
                      Burgers' equation dataset from
                      [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator)
+
+                     mapping between initial conditions to the solutions at the last point of time evolition in some function space. 
+                     u(x,0) -> u(x, time_end):
+
+                     * `a`: initial conditions u(x,0)
+                     * `u`: solutions u(x,t_end)
                      """,
+
                      "http://www.med.cgu.edu.tw/NeuralOperators/Burgers_R10.zip",
                      "9cbbe5070556c777b1ba3bacd49da5c36ea8ed138ba51b6ee76a24b971066ecd",
                      post_fetch_method = unpack))
@@ -63,6 +71,8 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 500)
                       ToDevice(device, device))
 
     fit!(learner, epochs)
+    model = learner.model |> cpu
+    @save "model/model_burger.bson" model
 
     return learner
 end
@@ -104,4 +114,12 @@ function train_nomad(; n = 300, cuda = true, learning_rate = 0.001, epochs = 400
     return mean_diff
 end
 
+
+function get_model()
+    model_path = joinpath(@__DIR__, "../model/")
+    model_file = readdir(model_path)[end]
+
+    return BSON.load(joinpath(model_path, model_file), @__MODULE__)[:model]
+end
+
 end

example/Burgers/src/Burgers.jl

@@ -8,7 +8,7 @@ using Test
     @test size(xs) == (2, 1024, 1000)
     @test size(ys) == (1, 1024, 1000)
 
-    learner = Burgers.train(epochs = 10)
+    learner = Burgers.train(epochs = 100)
     loss = learner.cbstate.metricsepoch[ValidationPhase()][:Loss].values[end]
     @test loss < 0.1