train MNO in 2-D and make sure loss<1e-2 about 5e-3 at epoch=35

Author: foldfelis

example/DoublePendulum/src/DoublePendulum.jl

@@ -15,7 +15,7 @@ function update_model!(model_file_path, model)
     @warn "model updated!"
 end
 
-function train(; loss_bounds=[1, 0.2, 0.1, 0.05, 0.02])
+function train(; loss_bounds=[0.05])
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -25,10 +25,10 @@ function train(; loss_bounds=[1, 0.2, 0.1, 0.05, 0.02])
     end
 
     m = Chain(
-        FourierOperator(6=>64, (16, ), relu),
-        FourierOperator(64=>64, (16, ), relu),
-        FourierOperator(64=>64, (16, ), relu),
-        FourierOperator(64=>6, (16, )),
+        FourierOperator(1=>64, (6, 64, ), relu),
+        FourierOperator(64=>64, (6, 64, ), relu),
+        FourierOperator(64=>64, (6, 64, ), relu),
+        FourierOperator(64=>1, (6, 64, )),
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]

example/DoublePendulum/src/data.jl

@@ -42,12 +42,12 @@ end
 
 function get_dataloader(; i=0, n_train=15001, n_test=1501, Δn=1024, batchsize=100)
     x = reshape(get_double_pendulum_chaotic_data(; i=i, n=-1), :)
-    𝐱 = reshape(vcat([x[i:(i+6Δn-1)] for i in 1:6:(length(x)-6(Δn-1))]...), 6, 1024, :)
+    𝐱 = reshape(vcat([x[i:(i+6Δn-1)] for i in 1:6:(length(x)-6(Δn-1))]...), 1, 6, 1024, :)
 
-    𝐱_train, 𝐲_train = 𝐱[:, :, 1:(n_train-1)], 𝐱[:, :, 2:n_train]
+    𝐱_train, 𝐲_train = 𝐱[:, :, :, 1:(n_train-1)], 𝐱[:, :, :, 2:n_train]
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
 
-    𝐱_test, 𝐲_test = 𝐱[:, :, (end-n_test+1):(end-1)], 𝐱[:, :, (end-n_test+2):end]
+    𝐱_test, 𝐲_test = 𝐱[:, :, :, (end-n_test+1):(end-1)], 𝐱[:, :, :, (end-n_test+2):end]
     loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
 
     return loader_train, loader_test