after NS model successful

Author: foldfelis

example/DoublePendulum/notebook/data.jl

@@ -14,37 +14,41 @@ begin
 end
 
 # ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
-data = DoublePendulum.get_data(i=0, n=-1)[:, end-2048+1:end]
+data, _, _, _ = DoublePendulum.preprocess(DoublePendulum.get_data(i=0), ratio=1);
+# data = reshape(DoublePendulum.get_data(; i=0, n=1000), 1, 1, 6, :);
 
 # ╔═╡ 4d0b08a4-8a54-41fd-997f-ad54d4c984cd
 m = DoublePendulum.get_model()
 
 # ╔═╡ 794374ce-6674-481d-8a3b-04db0f32d233
 begin
-	n = 10
+	n = 100
 
-	ground_truth_data = 1 .- data[:, 1+n:1024+n]
+	ground_truth_data = data[1, 1, :, 1:n]
 
-	inferenced_data = m(reshape(data[:, 1:1024], 1, 6, :, 1))
-	for i in 1:n
-		inferenced_data = m(inferenced_data)
-	end
-	inferenced_data = 1 .- reshape(inferenced_data, 6, :)
+	# inferenced_data = m(reshape(data[:, 1:1024], 1, 6, :, 1))
+	# for i in 1:n
+	# 	inferenced_data = m(inferenced_data)
+	# end
+	# inferenced_data = 1 .- reshape(inferenced_data, 6, :)
 end
 
 # ╔═╡ 9c8b3f8a-1b85-4c32-a416-ead51b244b94
 begin
-	anim = @animate for i in 1:4:1024
-		scatter(legend=false, xlim=(0, 1), ylim=(-0.5, 1), size=(600, 500))
-		scatter!(
-			inferenced_data[[2, 4, 6], i], inferenced_data[[1, 3, 5], i], 
-			color=[
-				RGB([239, 71, 111]/255...), 
-				RGB([6, 214, 160]/255...), 
-				RGB([17, 138, 178]/255...)
-			],
-			markersize=8
+	anim = @animate for i in 1:n
+		scatter(
+			legend=false, ticks=false, 
+			xlim=(0, 2500), ylim=(0, 2500), size=(600, 500)
 		)
+		# scatter!(
+		# 	inferenced_data[[2, 4, 6], i], inferenced_data[[1, 3, 5], i], 
+		# 	color=[
+		# 		RGB([239, 71, 111]/255...), 
+		# 		RGB([6, 214, 160]/255...), 
+		# 		RGB([17, 138, 178]/255...)
+		# 	],
+		# 	markersize=8
+		# )
 		scatter!(
 			ground_truth_data[[2, 4, 6], i], ground_truth_data[[1, 3, 5], i], 
 			color=[
@@ -54,7 +58,7 @@ begin
 			],
 			markersize=4
 		)
-		annotate!(0.1, -0.4, text("i=$i", :left))
+		annotate!(50, 50, text("t=$i", :left))
 	end
 
 	gif(anim)
@@ -66,4 +70,4 @@ end
 # ╠═5268feee-bda2-4612-9d4c-a1db424a11c7
 # ╠═4d0b08a4-8a54-41fd-997f-ad54d4c984cd
 # ╠═794374ce-6674-481d-8a3b-04db0f32d233
-# ╟─9c8b3f8a-1b85-4c32-a416-ead51b244b94
+# ╠═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(; loss_bounds=[])
+function train(; Δt=2)
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -25,33 +25,18 @@ function train(; loss_bounds=[])
     end
 
     m = Chain(
-        Dense(1, 64, gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, ), gelu),
-        FourierOperator(64=>64, (12, )),
-        Dense(64, 1)
+        Dense(1, 350), # (1, 2, 6, :) -> (350, 2, 6, :)
+        x -> reshape(x, 1, 60, 70, :), # (350, 2, 6, :) -> (1, 60, 70, :)
+        MarkovNeuralOperator(),
+        x -> reshape(x, 350, 2, 6, :), # (1, 60, 70, :) -> (350, 2, 6, :)
+        Dense(350, 1), # (350, 2, 6, :) -> (1, 2, 6, :)
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
 
     opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
 
-    loader_train, loader_test = get_dataloader()
-
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
+    loader_train, loader_test = get_dataloader(Δt=Δt)
 
     losses = Float32[]
     function validate()
@@ -60,16 +45,10 @@ function train(; loss_bounds=[])
 
         push!(losses, validation_loss)
         (losses[end] == minimum(losses)) && update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
-
-        isempty(loss_bounds) && return
-        if validation_loss < loss_bounds[1]
-            @warn "change η"
-            opt.os[2].eta /= 2
-            popfirst!(loss_bounds)
-        end
     end
     call_back = Flux.throttle(validate, 10, leading=false, trailing=true)
 
+    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
     Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
 end
 

example/DoublePendulum/src/data.jl

@@ -36,15 +36,25 @@ function get_data(; i=0, n=-1)
     return Float32.(data)
 end
 
-function get_dataloader(; i=0, n_train=15733, n_test=2048, Δn=1, batchsize=100)
-    𝐱 = get_data(i=i, n=-1) # size==(6, 17782)
-    ∇𝐱 = 𝐱[:, (1+Δn):end] - 𝐱[:, 1:(end-Δn)]
-    𝐱 = reshape(vcat(𝐱[:, 1:(end-Δn)], ∇𝐱), 1, 12, :)
+function preprocess(𝐱; Δt=2, ratio=0.9)
+    𝐱 = reshape(𝐱[:, 1:Δt:end], 1, 1, 6, :)
+    ∇𝐱 = 𝐱[:, :, :, 2:end] - 𝐱[:, :, :, 1:(end-1)]
 
-    𝐱_train, 𝐲_train = 𝐱[:, :, 1:(n_train-1)], 𝐱[:, :, 2:n_train]
-    loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
+    𝐱 = cat(𝐱[:, :, :, 1:(end-1)], ∇𝐱, dims=2)
+
+    n_train, n_test = floor(Int, ratio*size(𝐱)[end]), floor(Int, (1-ratio)*size(𝐱)[end])
 
-    𝐱_test, 𝐲_test = 𝐱[:, :, (end-n_test+1):(end-1)], 𝐱[:, :, (end-n_test+2):end]
+    𝐱_train, 𝐲_train = 𝐱[:, :, :, 1:(n_train-1)], 𝐱[:, :, :, 2:n_train]
+    𝐱_test, 𝐲_test = 𝐱[:, :, :, (end-n_test+1):(end-1)], 𝐱[:, :, :, (end-n_test+2):end]
+
+    return 𝐱_train, 𝐲_train, 𝐱_test, 𝐲_test
+end
+
+function get_dataloader(; i=0, Δt=2, ratio=0.9, batchsize=100)
+    𝐱 = get_data(i=i) # size==(6, :)
+    𝐱_train, 𝐲_train, 𝐱_test, 𝐲_test = preprocess(𝐱, Δt=Δt, ratio=ratio) # size==(1, 2, 6, :)
+
+    loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
     loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
 
     return loader_train, loader_test