Merge pull request #6 from foldfelis/example

Example for Burger's equation

Author: foldfelis

Project.toml

@@ -13,6 +13,13 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+DataDeps = "0.7"
+FFTW = "1.4"
+Fetch = "0.1"
+Flux = "0.12"
+MAT = "0.10"
+Tullio = "0.3"
+Zygote = "0.6"
 julia = "1.6"
 
 [extras]

example/burgers.jl

@@ -0,0 +1,36 @@
+using NeuralOperators
+using Flux
+# using CUDA
+
+# if has_cuda()
+#     @info "CUDA is on"
+#     device = gpu
+#     CUDA.allowscalar(false)
+# else
+    device = cpu
+# end
+
+m = FourierNeuralOperator() |> device
+loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+
+𝐱, 𝐲 = get_burgers_data(n=2048)
+
+n_train = 2000
+𝐱_train, 𝐲_train = 𝐱[:, :, 1:n_train], 𝐲[:, 1:n_train]
+loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=20, shuffle=true)
+
+n_test = 40
+𝐱_test, 𝐲_test = 𝐱[:, :, end-n_test+1:end], 𝐲[:, end-n_test+1:end]
+loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=20, shuffle=false)
+
+function loss_test()
+    l = 0
+    for (𝐱, 𝐲) in loader_test
+        l += loss(𝐱, 𝐲)
+    end
+    @info "loss: $(l/length(loader_test))"
+end
+
+data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
+opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=Flux.throttle(loss_test, 10)))

src/data.jl

@@ -15,13 +15,13 @@ function register_datasets()
     ))
 end
 
-function get_burgers_data(; n=1000, Δsamples=2^3, grid_size=div(2^13, Δsamples))
+function get_burgers_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples), T=Float32)
     file = matopen(joinpath(datadep"BurgersR10", "burgers_data_R10.mat"))
-    x_data = collect(read(file, "a")[1:n, 1:Δsamples:end]')
-    y_data = collect(read(file, "u")[1:n, 1:Δsamples:end]')
+    x_data = T.(collect(read(file, "a")[1:n, 1:Δsamples:end]'))
+    y_data = T.(collect(read(file, "u")[1:n, 1:Δsamples:end]'))
     close(file)
 
-    x_loc_data = Array{Float32, 3}(undef, 2, grid_size, n)
+    x_loc_data = Array{T, 3}(undef, 2, grid_size, n)
     x_loc_data[1, :, :] .= reshape(repeat(LinRange(0, 1, grid_size), n), (grid_size, n))
     x_loc_data[2, :, :] .= x_data
 

test/data.jl

@@ -1,5 +1,5 @@
 @testset "get burgers data" begin
-    xs, ys = get_burgers_data()
+    xs, ys = get_burgers_data(n=1000)
 
     @test size(xs) == (2, 1024, 1000)
     @test size(ys) == (1024, 1000)

test/fourier.jl

@@ -7,7 +7,7 @@
         SpectralConv1d(ch, modes)
     )
 
-    𝐱, _ = get_burgers_data()
+    𝐱, _ = get_burgers_data(n=1000)
     @test size(m(𝐱)) == (64, 1024, 1000)
 
     T = Float32
@@ -25,7 +25,7 @@ end
         FourierOperator(ch, modes)
     )
 
-    𝐱, _ = get_burgers_data()
+    𝐱, _ = get_burgers_data(n=1000)
     @test size(m(𝐱)) == (64, 1024, 1000)
 
     loss(x, y) = Flux.mse(m(x), y)