revise data

Author: foldfelis

example/DoublePendulum/notebook/data.jl

@@ -0,0 +1,88 @@
+### A Pluto.jl notebook ###
+# v0.15.1
+
+using Markdown
+using InteractiveUtils
+
+# ╔═╡ 194baef2-0417-11ec-05ab-4527ef614024
+using Pkg; Pkg.develop(path=".."); Pkg.activate("..")
+
+# ╔═╡ 38c9ced5-dcf8-4e03-ac07-7c435687861b
+begin
+	using DoublePendulum
+	using Plots
+end
+
+# ╔═╡ 396b5d7a-a7a4-4f22-a87e-39b405e8d62a
+md"
+# Double Pendulum
+
+JingYu Ning
+"
+
+# ╔═╡ 2a606ecf-acf0-41ad-9290-7569dbb22b5a
+md"
+The data is provided by [IBM](https://developer.ibm.com/exchanges/data/all/double-pendulum-chaotic/)
+
+> In this dataset, videos of the double pendulum were taken using a high-speed Phantom Miro EX2 camera. To make the extraction of the arm positions easier, a matte black background was used, and the three datums were marked with red, green and blue fiducial markers. The camera was placed at 2 meters from the pendulum, with the axis of the objective aligned with the first pendulum datum. The pendulum was launched by hand, and the camera was motion triggered. The dataset was generated on the basis of 21 individual runs of the pendulum. Each of the recorded sequences lasted around 40s and consisted of around 17500 frames.
+"
+
+# ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
+data, _, _, _ = DoublePendulum.preprocess(
+	DoublePendulum.get_data(i=10, n=410),
+	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]
+	
+	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])
+	end
+	inferenced_data = inferenced_data[1, :, :]
+end;
+
+# ╔═╡ 9c8b3f8a-1b85-4c32-a416-ead51b244b94
+begin
+	c = [
+		RGB([239, 71, 111]/255...), 
+		RGB([6, 214, 160]/255...), 
+		RGB([17, 138, 178]/255...)
+	]
+	xi, yi = [2, 4, 6], [1, 3, 5]
+	
+	anim = @animate for i in 1:n
+		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)
+		)
+		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))
+	end
+
+	gif(anim, fps=5)
+end
+
+# ╔═╡ Cell order:
+# ╟─396b5d7a-a7a4-4f22-a87e-39b405e8d62a
+# ╟─2a606ecf-acf0-41ad-9290-7569dbb22b5a
+# ╟─194baef2-0417-11ec-05ab-4527ef614024
+# ╠═38c9ced5-dcf8-4e03-ac07-7c435687861b
+# ╠═5268feee-bda2-4612-9d4c-a1db424a11c7
+# ╠═4d0b08a4-8a54-41fd-997f-ad54d4c984cd
+# ╠═794374ce-6674-481d-8a3b-04db0f32d233
+# ╟─9c8b3f8a-1b85-4c32-a416-ead51b244b94

example/DoublePendulum/notebook/double_pendulum.jl

@@ -25,11 +25,11 @@ function train(; Δt=2)
     end
 
     m = Chain(
-        Dense(1, 350), # (1, 2, 6, :) -> (350, 2, 6, :)
-        x -> reshape(x, 1, 60, 70, :), # (350, 2, 6, :) -> (1, 60, 70, :)
+        Dense(2, Int(4096/4)),
+        x -> reshape(x, 1, 64, 64, :),
         MarkovNeuralOperator(),
-        x -> reshape(x, 350, 2, 6, :), # (1, 60, 70, :) -> (350, 2, 6, :)
-        Dense(350, 1), # (350, 2, 6, :) -> (1, 2, 6, :)
+        x -> reshape(x, Int(4096/4), 4, :),
+        Dense(Int(4096/4), 2),
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]

example/DoublePendulum/src/DoublePendulum.jl

@@ -37,22 +37,32 @@ function get_data(; i=0, n=-1)
 end
 
 function preprocess(𝐱; Δt=2, ratio=0.9)
-    𝐱 = reshape(𝐱[:, 1:Δt:end], 1, 1, 6, :)
-    ∇𝐱 = 𝐱[:, :, :, 2:end] - 𝐱[:, :, :, 1:(end-1)]
+    xs_red, ys_red = 𝐱[1, :], 𝐱[2, :]
+    𝐱[1, :] -= xs_red; 𝐱[3, :] -= xs_red; 𝐱[5, :] -= xs_red
+    𝐱[2, :] -= ys_red; 𝐱[4, :] -= ys_red; 𝐱[6, :] -= ys_red
 
-    𝐱 = cat(𝐱[:, :, :, 1:(end-1)], ∇𝐱, dims=2)
+    𝐱 = reshape(𝐱[3:6, 1:Δt:end], 1, 4, :)
+    ∇𝐱 = 𝐱[:, :, 2:end] - 𝐱[:, :, 1:(end-1)]
+
+    𝐱 = cat(𝐱[:, :, 1:(end-1)], ∇𝐱, dims=1)
 
     n_train, n_test = floor(Int, ratio*size(𝐱)[end]), floor(Int, (1-ratio)*size(𝐱)[end])
 
-    𝐱_train, 𝐲_train = 𝐱[:, :, :, 1:(n_train-1)], 𝐱[:, :, :, 2:n_train]
-    𝐱_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, :)
+function get_dataloader(; n_file=10, Δ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)
+    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, 𝐲_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)
+    end
 
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
     loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)