Transferred DeepONet training task to example

Author: ba2tripleO

Project.toml

@@ -10,7 +10,6 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
-MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

example/Burgers/src/Burgers.jl

@@ -5,6 +5,7 @@ using Flux
 using CUDA
 
 include("data.jl")
+include("Burgers_deeponet.jl")
 
 __init__() = register_burgers()
 
@@ -30,7 +31,7 @@ function train()
         Dense(128, 1),
         flatten
     ) |> device
-    
+
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
 
     loader_train, loader_test = get_dataloader()

example/Burgers/src/Burgers_deeponet.jl

@@ -0,0 +1,32 @@
+function train_don()
+    if has_cuda()
+        @info "CUDA is on"
+        device = gpu
+        CUDA.allowscalar(false)
+    else
+        device = cpu
+    end
+
+    x, y = get_data_don(n=300)
+    xtrain = x[1:280, :]' |> device
+    xval = x[end-19:end, :]' |device
+
+    ytrain = y[1:280, :] |> device
+    yval = y[end-19:end, :] |> device
+
+    grid = collect(range(0, 1, length=1024))' |> device
+
+    learning_rate = 0.001
+    opt = ADAM(learning_rate)
+
+    m = DeepONet((1024,1024,1024),(1,1024,1024),gelu,gelu)
+    loss(xtrain,ytrain,sensor) = Flux.Losses.mse(model(xtrain,sensor),ytrain)
+    evalcb() = @show(loss(xval,yval,grid))
+
+    Flux.@epochs 400 Flux.train!(loss, params(m), [(xtrain,ytrain,grid)], opt, cb = evalcb)
+    ỹ = m(xval, grid)
+
+    diffvec = vec(abs.((yval .- ỹ)))
+    mean_diff = sum(diffvec)/length(diffvec)
+    return mean_diff
+end

example/Burgers/src/data.jl

@@ -27,6 +27,15 @@ function get_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples), T=Flo
     return x_loc_data, y_data
 end
 
+function get_data_don(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples))
+    file = matopen(joinpath(datadep"Burgers", "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])
+    close(file)
+
+    return x_data, y_data
+end
+
 function get_dataloader(; n_train=1800, n_test=200, batchsize=100)
     𝐱, 𝐲 = get_data(n=2048)
 

example/Burgers/test/deeponet.jl

@@ -0,0 +1,5 @@
+@testset "DeepONet Training Accuracy" begin
+    ϵ = Burgers.train_don()
+
+    @test ϵ < 0.4
+end

example/Burgers/test/runtests.jl

@@ -3,4 +3,5 @@ using Test
 
 @testset "Burgers" begin
     include("data.jl")
+    include("deeponet.jl")
 end

test/burgerset.mat

@@ -34,34 +34,3 @@ mgrad = Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)
 @test !iszero(Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)[2])
 
 #training
-#dataset containing first 300 initial conditions from the Burgers' equation
-#dataset used by Li et al. for Fourier neural operator. The data for the initial
-#conditions is sampled at an interval of 8 points, so, the original data has
-#2048 ICs at 8192 points, while here we have 300 ICs at 1024 points
-vars = matread("burgerset.mat")
-
-xtrain = vars["a"][1:280, :]'
-xval = vars["a"][end-19:end, :]'
-
-ytrain = vars["u"][1:280, :]
-yval = vars["u"][end-19:end, :]
-
-grid = collect(range(0, 1, length=1024))'
-model = DeepONet((1024,1024,1024),(1,1024,1024),gelu,gelu)
-
-learning_rate = 0.001
-opt = ADAM(learning_rate)
-
-parameters = params(model)
-
-loss(xtrain,ytrain,sensor) = Flux.Losses.mse(model(xtrain,sensor),ytrain)
-
-evalcb() = @show(loss(xval,yval,grid))
-
-Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain,ytrain,grid)], opt, cb = evalcb)
-
-ỹ = model(xval, grid)
-
-diffvec = vec(abs.((yval .- ỹ)))
-mean_diff = sum(diffvec)/length(diffvec)
-@test mean_diff < 0.4