use all data

Author: foldfelis

example/DoublePendulum/notebook/double_pendulum.jl

@@ -29,7 +29,7 @@ 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=10, n=410),
+	DoublePendulum.get_data(i=20, n=410),
 	ratio=1
 );
 
@@ -39,9 +39,9 @@ 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
@@ -53,18 +53,18 @@ end;
 # ╔═╡ 9c8b3f8a-1b85-4c32-a416-ead51b244b94
 begin
 	c = [
-		RGB([239, 71, 111]/255...), 
-		RGB([6, 214, 160]/255...), 
+		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, 
+			legend=false, ticks=false,
 			xlim=(-1000, 1000), ylim=(-1000, 1000), size=(400, 350)
 		)
 		plot!(i_data[xi], i_data[yi], color=:black)

example/DoublePendulum/src/DoublePendulum.jl

@@ -49,7 +49,11 @@ function train(; Δt=2)
     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))
+    for e in 1:50
+        @info "Epoch $e"
+        @time Flux.train!(loss, params(m), data, opt, cb=call_back)
+        (e%3 == 0) && (opt.os[2].eta /= 2)
+    end
 end
 
 function get_model()

example/DoublePendulum/src/data.jl

@@ -37,11 +37,14 @@ function get_data(; i=0, n=-1)
 end
 
 function preprocess(𝐱; Δt=2, ratio=0.9)
+    # move red point to (0, 0)
     xs_red, ys_red = 𝐱[1, :], 𝐱[2, :]
-    𝐱[1, :] -= xs_red; 𝐱[3, :] -= xs_red; 𝐱[5, :] -= xs_red
-    𝐱[2, :] -= ys_red; 𝐱[4, :] -= ys_red; 𝐱[6, :] -= ys_red
+    𝐱[3, :] -= xs_red; 𝐱[5, :] -= xs_red
+    𝐱[4, :] -= ys_red; 𝐱[6, :] -= ys_red
 
+    # needs only green and blue points
     𝐱 = reshape(𝐱[3:6, 1:Δt:end], 1, 4, :)
+    # velocity of green and blue points
     ∇𝐱 = 𝐱[:, :, 2:end] - 𝐱[:, :, 1:(end-1)]
 
     𝐱 = cat(𝐱[:, :, 1:(end-1)], ∇𝐱, dims=1)
@@ -54,7 +57,7 @@ function preprocess(𝐱; Δt=2, ratio=0.9)
     return 𝐱_train, 𝐲_train, 𝐱_test, 𝐲_test
 end
 
-function get_dataloader(; n_file=10, Δt=2, ratio=0.9, batchsize=100)
+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)
     for i in 0:(n_file-1)