implement super resolution mno

Author: foldfelis

example/FlowOverCircle/src/data.jl

@@ -22,7 +22,7 @@ function gen_data(ts::AbstractRange)
     𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
     for (i, t) in enumerate(ts)
         sim_step!(circ, t)
-        𝐩s[:, :, :, i] = Float32.(circ.flow.p)[2:end-1, 2:end-1]
+        𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
     end
 
     return 𝐩s

example/SuperResolution/README.md

@@ -1,16 +1,18 @@
 # Super Resolution
 
 The time dependent Navier-Stokes equation is learned by the `MarkovNeuralOperator` with only one time step information.
-The result of this example can be found [here](https://foldfelis.github.io/NeuralOperators.jl/dev/assets/notebook/mno.jl.html).
+The result of this example can be found [here](https://foldfelis.github.io/NeuralOperators.jl/dev/assets/notebook/super_resolution_mno.jl.html).
+
+Apart from just training a MNO, here, we train the model with lower resolution (96x64) and inference result with higher resolution (192x128).
 
 | **Ground Truth** | **Inferenced** |
 |:----------------:|:--------------:|
 | ![](gallery/ans.gif) | ![](gallery/inferenced.gif) |
 
-Change directory to `example/FlowOverCircle` and use following commend to train model:
+Change directory to `example/SuperResolution` and use following commend to train model:
 
 ```julia
 $ julia --proj
 
-julia> using FlowOverCircle; FlowOverCircle.train()
+julia> using SuperResolution; SuperResolution.train()
 ```

example/SuperResolution/gallery/ans.gif

@@ -1,5 +1,5 @@
 ### A Pluto.jl notebook ###
-# v0.15.1
+# v0.16.1
 
 using Markdown
 using InteractiveUtils
@@ -8,11 +8,11 @@ using InteractiveUtils
 using Pkg; Pkg.develop(path=".."); Pkg.activate("..")
 
 # ╔═╡ 38c9ced5-dcf8-4e03-ac07-7c435687861b
-using FlowOverCircle, Plots
+using SuperResolution, Plots
 
 # ╔═╡ 50ce80a3-a1e8-4ba9-a032-dad315bcb432
 md"
-# Markov Neural Operator
+# Super Resolution with MNO
 
 JingYu Ning
 "
@@ -28,16 +28,26 @@ md"
 
 # ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
 begin 
-	n = 20
-	data = FlowOverCircle.gen_data(LinRange(100, 100+n-1, n))
+	n = 10
+	data = SuperResolution.gen_data(LinRange(100, 100+n-1, n))
 end;
 
-# ╔═╡ 9b02b6a2-33c3-4ca6-bfba-0bd74b664830
+# ╔═╡ 5531bba6-94bd-4c99-be8c-43fe19ad8a60
+md"
+## Training
+"
+
+# ╔═╡ 74fc528f-ccd4-4670-9b17-dbfa7a1c74b6
+md"
+Apart from just training a MNO, here, we train the model with lower resolution (96x64) and inference result with higher resolution (192x128). 
+"
+
+# ╔═╡ f6d1ce85-a195-4ab1-bd3a-dbd4b0d1fcca
 begin
 	anim = @animate for i in 1:size(data)[end]
-		heatmap(data[1, :, :, i]', color=:coolwarm, clim=(-1.5, 1.5))
+		heatmap(data[1, 1:2:end, 1:2:end, i]', color=:coolwarm, clim=(-1.5, 1.5))
 		scatter!(
-			[size(data, 3)÷2], [size(data, 3)÷2-1], 
+			[size(data, 3)÷4-1], [size(data, 3)÷4-1], 
 			markersize=45, color=:black, legend=false, ticks=false
 		)
 		annotate!(5, 5, text("i=$i", :left))
@@ -54,7 +64,7 @@ Use the first data generated above as the initial state, and apply the operator
 
 # ╔═╡ fbc287b8-f232-4350-9948-2091908e5a30
 begin
-	m = FlowOverCircle.get_model()
+	m = SuperResolution.get_model()
 
 	states = Array{Float32}(undef, size(data))
 	states[:, :, :, 1] .= view(data, :, :, :, 1)
@@ -68,7 +78,7 @@ begin
 	anim_model = @animate for i in 1:size(states)[end]
 		heatmap(states[1, :, :, i]', color=:coolwarm, clim=(-1.5, 1.5))
 		scatter!(
-			[size(data, 3)÷2], [size(data, 3)÷2-1], 
+			[size(data, 3)÷2-1], [size(data, 3)÷2-1], 
 			markersize=45, color=:black, legend=false, ticks=false
 		)
 		annotate!(5, 5, text("i=$i", :left))
@@ -79,11 +89,13 @@ end
 # ╔═╡ Cell order:
 # ╟─50ce80a3-a1e8-4ba9-a032-dad315bcb432
 # ╟─59769504-ebd5-4c6f-981f-d03826d8e34a
-# ╟─194baef2-0417-11ec-05ab-4527ef614024
+# ╠═194baef2-0417-11ec-05ab-4527ef614024
 # ╠═38c9ced5-dcf8-4e03-ac07-7c435687861b
 # ╟─823b3547-6723-43cf-85e6-cc6eb44efea1
 # ╠═5268feee-bda2-4612-9d4c-a1db424a11c7
-# ╟─9b02b6a2-33c3-4ca6-bfba-0bd74b664830
+# ╟─5531bba6-94bd-4c99-be8c-43fe19ad8a60
+# ╟─74fc528f-ccd4-4670-9b17-dbfa7a1c74b6
+# ╠═f6d1ce85-a195-4ab1-bd3a-dbd4b0d1fcca
 # ╟─55058635-c7e9-4ee3-81c2-0153e84f4c8e
 # ╠═fbc287b8-f232-4350-9948-2091908e5a30
 # ╟─a0b5e94c-a839-4cc0-a325-1a4ac39fafbc

example/SuperResolution/gallery/inferenced.gif

@@ -16,13 +16,13 @@ function circle(n, m; Re=250)
 end
 
 function gen_data(ts::AbstractRange)
-    n, m = 3(2^5), 2^6
+    n, m = 2 * 3(2^5), 2 * 2^6
     circ = circle(n, m)
 
     𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
     for (i, t) in enumerate(ts)
         sim_step!(circ, t)
-        𝐩s[:, :, :, i] = Float32.(circ.flow.p)[2:end-1, 2:end-1]
+        𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
     end
 
     return 𝐩s
@@ -33,7 +33,7 @@ function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::
 
     n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))
 
-    𝐱_train, 𝐲_train = data[:, :, :, 1:(n_train-1)], data[:, :, :, 2:n_train]
+    𝐱_train, 𝐲_train = data[:, 1:2:end, 1:2:end, 1:(n_train-1)], data[:, 1:2:end, 1:2:end, 2:n_train]
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
 
     𝐱_test, 𝐲_test = data[:, :, :, (end-n_test+1):(end-1)], data[:, :, :, (end-n_test+2):end]