train task via 2-D NFO

Author: foldfelis

example/DoublePendulum/notebook/double_pendulum.jl

@@ -1,5 +1,5 @@
 ### A Pluto.jl notebook ###
-# v0.15.1
+# v0.16.0
 
 using Markdown
 using InteractiveUtils
@@ -28,26 +28,35 @@ The data is provided by [IBM](https://developer.ibm.com/exchanges/data/all/doubl
 "
 
 # ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
-data, _, _, _ = DoublePendulum.preprocess(
-	DoublePendulum.get_data(i=20, n=410),
+data_x, data_y, _, _ = DoublePendulum.preprocess(
+	DoublePendulum.get_data(i=20),
 	ratio=1
 );
 
 # ╔═╡ 4d0b08a4-8a54-41fd-997f-ad54d4c984cd
 m = DoublePendulum.get_model();
 
-# ╔═╡ 794374ce-6674-481d-8a3b-04db0f32d233
-begin
-	n = 20
-
-	ground_truth_data = data[1, :, 1:n]
+# ╔═╡ ad6302b2-3d62-4a3f-b8bf-f69bab80c7a4
+ground_truth_data = cat(
+	[data_x[:, :, :, i] 
+	for i in 1:size(data_x, 3):size(data_x)[end]]..., dims=3
+)[1, :, :];
 
-	inferenced_data = Array{Float32}(undef, 2, 4, n)
-	inferenced_data[:, :, 1] .= data[:, :, 1]
-	for i in 2:n
-		inferenced_data[:, :, i:i] .= m(inferenced_data[:, :, i-1:i-1])
+# ╔═╡ 794374ce-6674-481d-8a3b-04db0f32d233
+begin 
+	n = 5
+	inferenced_data = data_x[:, :, :, 1:1]
+	for i in 1:n
+		inferenced_data = cat(
+			inferenced_data, 
+			m(inferenced_data[:, :, :, i:i]), 
+			dims=4
+		)
 	end
-	inferenced_data = inferenced_data[1, :, :]
+	
+	inferenced_data = cat(
+		[inferenced_data[:, :, :, i] for i in 1:n]..., dims=3
+	)[1, :, :]
 end;
 
 # ╔═╡ 9c8b3f8a-1b85-4c32-a416-ead51b244b94
@@ -59,22 +68,27 @@ begin
 	]
 	xi, yi = [2, 4, 6], [1, 3, 5]
 
-	anim = @animate for i in 1:n
+	anim = @animate for i in 1:size(inferenced_data)[end]
 		i_data = [0, 0, inferenced_data[:, i]...]
 		g_data = [0, 0, ground_truth_data[:, i]...]
 
 		scatter(
 			legend=false, ticks=false,
-			xlim=(-1000, 1000), ylim=(-1000, 1000), size=(400, 350)
+			xlim=(-1500, 1500), ylim=(-1500, 1500), size=(400, 350)
 		)
 		plot!(i_data[xi], i_data[yi], color=:black)
 		scatter!(i_data[xi], i_data[yi], color=c, markersize=8)
 		plot!(g_data[xi], g_data[yi], color=:gray)
 		scatter!(g_data[xi], g_data[yi], color=c, markersize=4)
-		annotate!(-900, -900, text("t=$i", :left))
+		
+		if i ≤ 30
+			annotate!(-1400, -1400, text("t=$i", :left, color=:black))
+		else
+			annotate!(-1400, -1400, text("t=$i", :left, color=:red))
+		end
 	end
 
-	gif(anim, fps=5)
+	gif(anim)
 end
 
 # ╔═╡ Cell order:
@@ -84,5 +98,6 @@ end
 # ╠═38c9ced5-dcf8-4e03-ac07-7c435687861b
 # ╠═5268feee-bda2-4612-9d4c-a1db424a11c7
 # ╠═4d0b08a4-8a54-41fd-997f-ad54d4c984cd
+# ╠═ad6302b2-3d62-4a3f-b8bf-f69bab80c7a4
 # ╠═794374ce-6674-481d-8a3b-04db0f32d233
 # ╟─9c8b3f8a-1b85-4c32-a416-ead51b244b94

example/DoublePendulum/src/DoublePendulum.jl

@@ -15,7 +15,7 @@ function update_model!(model_file_path, model)
     @warn "model updated!"
 end
 
-function train(; Δt=2)
+function train(; Δt=1)
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -25,11 +25,13 @@ function train(; Δt=2)
     end
 
     m = Chain(
-        Dense(2, Int(4096/4)),
-        x -> reshape(x, 1, 64, 64, :),
-        MarkovNeuralOperator(),
-        x -> reshape(x, Int(4096/4), 4, :),
-        Dense(Int(4096/4), 2),
+        Dense(2, 64),
+        FourierOperator(64=>64, (4, 16), gelu),
+        FourierOperator(64=>64, (4, 16), gelu),
+        FourierOperator(64=>64, (4, 16), gelu),
+        FourierOperator(64=>64, (4, 16)),
+        Dense(64, 128, gelu),
+        Dense(128, 2),
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
@@ -49,8 +51,8 @@ function train(; Δt=2)
     call_back = Flux.throttle(validate, 10, leading=false, trailing=true)
 
     data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    for e in 1:50
-        @info "Epoch $e"
+    for e in 1:20
+        @info "Epoch $e\n η: $(opt.os[2].eta)"
         @time Flux.train!(loss, params(m), data, opt, cb=call_back)
         (e%3 == 0) && (opt.os[2].eta /= 2)
     end

example/DoublePendulum/src/data.jl

@@ -36,7 +36,7 @@ function get_data(; i=0, n=-1)
     return Float32.(data)
 end
 
-function preprocess(𝐱; Δt=2, ratio=0.9)
+function preprocess(𝐱; Δt=1, nx=30, ny=30, ratio=0.9)
     # move red point to (0, 0)
     xs_red, ys_red = 𝐱[1, :], 𝐱[2, :]
     𝐱[3, :] -= xs_red; 𝐱[5, :] -= xs_red
@@ -46,25 +46,33 @@ function preprocess(𝐱; Δt=2, ratio=0.9)
     𝐱 = reshape(𝐱[3:6, 1:Δt:end], 1, 4, :)
     # velocity of green and blue points
     ∇𝐱 = 𝐱[:, :, 2:end] - 𝐱[:, :, 1:(end-1)]
-
+    # merge info of pos and velocity
     𝐱 = cat(𝐱[:, :, 1:(end-1)], ∇𝐱, dims=1)
 
-    n_train, n_test = floor(Int, ratio*size(𝐱)[end]), floor(Int, (1-ratio)*size(𝐱)[end])
+    # with info of first nx steps to inference next ny steps
+    n = size(𝐱)[end] - (nx + ny) + 1
+    𝐱s = Array{Float32}(undef, size(𝐱)[1:2]..., nx, n)
+    𝐲s = Array{Float32}(undef, size(𝐱)[1:2]..., ny, n)
+    for i in 1:n
+        𝐱s[:, :, :, i] .= 𝐱[:, :, i:(i+nx-1)]
+        𝐲s[:, :, :, i] .= 𝐱[:, :, (i+nx):(i+nx+ny-1)]
+    end
 
-    𝐱_train, 𝐲_train = 𝐱[:, :, 1:(n_train-1)], 𝐱[:, :, 2:n_train]
-    𝐱_test, 𝐲_test = 𝐱[:, :, (end-n_test+1):(end-1)], 𝐱[:, :, (end-n_test+2):end]
+    n_train = floor(Int, ratio*n)
+    𝐱_train, 𝐲_train = 𝐱s[:, :, :, 1:n_train], 𝐲s[:, :, :, 1:n_train]
+    𝐱_test, 𝐲_test = 𝐱s[:, :, :, (n_train+1):end], 𝐲s[:, :, :, (n_train+1):end]
 
     return 𝐱_train, 𝐲_train, 𝐱_test, 𝐲_test
 end
 
-function get_dataloader(; n_file=20, Δt=2, ratio=0.9, batchsize=100)
-    𝐱_train, 𝐲_train = Array{Float32}(undef, 2, 4, 0), Array{Float32}(undef, 2, 4, 0)
-    𝐱_test, 𝐲_test = Array{Float32}(undef, 2, 4, 0), Array{Float32}(undef, 2, 4, 0)
+function get_dataloader(; n_file=20, Δt=1, nx=30, ny=30, ratio=0.9, batchsize=100)
+    𝐱_train, 𝐲_train = Array{Float32}(undef, 2, 4, nx, 0), Array{Float32}(undef, 2, 4, ny, 0)
+    𝐱_test, 𝐲_test = Array{Float32}(undef, 2, 4, nx, 0), Array{Float32}(undef, 2, 4, ny, 0)
     for i in 0:(n_file-1)
-        𝐱_train_i, 𝐲_train_i, 𝐱_test_i, 𝐲_test_i = preprocess(get_data(i=i), Δt=Δt, ratio=ratio)
+        𝐱_train_i, 𝐲_train_i, 𝐱_test_i, 𝐲_test_i = preprocess(get_data(i=i), Δt=Δt, nx=nx, ny=ny, ratio=ratio)
 
-        𝐱_train, 𝐲_train = cat(𝐱_train, 𝐱_train_i, dims=3), cat(𝐲_train, 𝐲_train_i, dims=3)
-        𝐱_test, 𝐲_test = cat(𝐱_test, 𝐱_test_i, dims=3), cat(𝐲_test, 𝐲_test_i, dims=3)
+        𝐱_train, 𝐲_train = cat(𝐱_train, 𝐱_train_i, dims=4), cat(𝐲_train, 𝐲_train_i, dims=4)
+        𝐱_test, 𝐲_test = cat(𝐱_test, 𝐱_test_i, dims=4), cat(𝐲_test, 𝐲_test_i, dims=4)
     end
 
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)