Optimize NS rhs, finalize get_data

.gitignore

@@ -4,3 +4,6 @@ examples/test*
 Manifest.toml
 /docs/Manifest.toml
 /docs/build/
+
+# Ignore large data files
+simulations/*/data/*.jld2

Project.toml

@@ -49,6 +49,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 ParameterSchedulers = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"
 Plotly = "58dd65bb-95f3-509e-9936-c39a10fdeae7"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"

simulations/NavierStokes_2D/scripts/get_data.jl

@@ -9,6 +9,7 @@ using Random
 using DifferentialEquations
 using Zygote
 using JLD2
+using PreallocationTools
 const ArrayType = Array
 const solver_algo = Tsit5()
 const MY_TYPE = Float32 # use float32 if you plan to use a GPU
@@ -32,14 +33,14 @@ CUDA.allowscalar(false)
 
 ## *** Generate the data with the following parameters
 nu = 5.0f-4
-les_size = 64
-dns_size = 128
+#les_size = 64
+#dns_size = 128
 les_size = 32
 dns_size = 64
 myseed = 1234
 data_name = get_data_name(nu, les_size, dns_size, myseed)
 # If the data already exists, we stop here
-if isfile("../data/$(data_name).jld2")
+if isfile("./simulations/NavierStokes_2D/data/$(data_name).jld2")
     println("Data already exists, stopping here.")
     return
 end
@@ -55,7 +56,14 @@ params_dns = create_params(dns_size; nu)
 Random.seed!(myseed)
 
 ## Initial conditions
-u0_dns = random_field(params_dns)
+nsamp = 50
+u0_dns = random_field(params_dns, nsamp = nsamp)
+# and we make sure that the initial condition is divergence free
+maximum(abs, params_dns.k .* u0_dns[:, :, 1, :] .+ params_dns.k' .* u0_dns[:, :, 2, :])
+
+# Create that cache which is used to compute the right hand side of the Navier-Stokes
+cache = create_cache(u0_dns);
+cache_les = create_cache(spectral_cutoff(u0_dns, params_les.K));
 
 ## Let's do some time stepping.
 dt_burn = 2.0f-4
@@ -67,7 +75,7 @@ trange_burn = (0.0f0, nburn * dt_burn)
 # run the burnout 
 import DiffEqFlux: NeuralODE
 include("./../../../src/NODE.jl")
-F_dns(u) = Zygote.@ignore project(F_NS(u, params_dns), params_dns)
+F_dns(u) = Zygote.@ignore project(F_NS(u, params_dns, cache), params_dns, cache)
 f_dns = create_f_CNODE((F_dns,); is_closed = false);
 import Random, Lux
 Random.seed!(123)
@@ -76,17 +84,17 @@ rng = Random.default_rng()
 dns_burn = NeuralODE(f_dns,
     trange_burn,
     solver_algo,
+    p = LazyBufferCache(),
     adaptive = false,
     dt = dt_burn,
     saveat = saveat_burn);
-u_dns_burn = dns_burn(u0_dns, θ_dns, st_dns)[1]
-u_dns_burn.u
+@time u_dns_burn = dns_burn(u0_dns, θ_dns, st_dns)[1];
 
-# and plot the burnout
+# and plot the burnout of one solution
 if plotting
     plot()
     @gif for (idx, (t, u)) in enumerate(zip(u_dns_burn.t, u_dns_burn.u))
-        ω = Array(vorticity(u, params_dns))
+        ω = Array(vorticity(u[:, :, :, 1], params_dns))
         title1 = @sprintf("Vorticity, burnout, t = %.3f", t)
         p1 = heatmap(ω'; xlabel = "x", ylabel = "y", titlefontsize = 11, title = title1)
         plot!(p1)
@@ -98,6 +106,9 @@ u0_dns_eq = u_dns_burn.u[end]
 u_dns_burn = nothing
 GC.gc()
 
+include("./../../../src/NavierStokes.jl")
+spectral_cutoff(u0_dns_eq, params_les.K)
+
 # Now we run the DNS and we compute the LES information at every time step
 dt = 2.0f-4
 nt = 1000
@@ -110,38 +121,29 @@ dns = NeuralODE(f_dns,
     saveat = dt);
 sol = dns(u0_dns_eq, θ_dns, st_dns)[1]
 # Preallocate memory for the LES data and the commutator
-v = zeros(Complex{MY_TYPE}, params_les.N, params_les.N, 2, nt + 1)
-c = zeros(Complex{MY_TYPE}, params_les.N, params_les.N, 2, nt + 1)
+v = zeros(Complex{MY_TYPE}, params_les.N, params_les.N, 2, nsamp, nt + 1)
+c = zeros(Complex{MY_TYPE}, params_les.N, params_les.N, 2, nsamp, nt + 1)
 
 # then we loop to plot and compute LES
 if plotting
     anim = Animation()
-    for (idx, (t, u)) in enumerate(zip(sol.t, sol.u))
-        ubar = spectral_cutoff(u, params_les.K)
-        v[:, :, :, idx] = Array(ubar)
-        c[:, :, :, idx] = Array(spectral_cutoff(F_NS(u, params_dns), params_les.K) -
-                                F_NS(ubar, params_les))
-        if idx % 10 == 0
-            ω = Array(vorticity(u, params_dns))
-            title = @sprintf("Vorticity, t = %.3f", t)
-            fig = heatmap(ω'; xlabel = "x", ylabel = "y", title)
-            frame(anim, fig)
-        end
-    end
-    gif(anim)
-else
-    for (idx, (t, u)) in enumerate(zip(sol.t, sol.u))
-        ubar = spectral_cutoff(u, params_les.K)
-        v[:, :, :, idx] = Array(ubar)
-        c[:, :, :, idx] = Array(spectral_cutoff(F(u, params_dns), params_les.K) -
-                                F(ubar, params_les))
+end
+for (idx, (t, u)) in enumerate(zip(sol.t, sol.u))
+    ubar = spectral_cutoff(u, params_les.K)
+    v[:, :, :, :, idx] = Array(ubar)
+    c[:, :, :, :, idx] = Array(spectral_cutoff(F_NS(u, params_dns, cache), params_les.K) -
+                               F_NS(ubar, params_les, cache_les))
+    if plotting && (idx % 10 == 0)
+        ω = Array(vorticity(u[:, :, :, 1], params_dns))
+        title = @sprintf("Vorticity, t = %.3f", t)
+        fig = heatmap(ω'; xlabel = "x", ylabel = "y", title)
+        frame(anim, fig)
     end
 end
-
-# create the directory data if it does not exist
-if !isdir("../data")
-    mkdir("../data")
+if plotting
+    gif(anim)
 end
 
 # Save all the simulation data
-save("../data/$(data_name).jld2", "data", Data(sol.t, sol.u, v, c, params_les, params_dns))
+save("./simulations/NavierStokes_2D/data/$(data_name).jld2",
+    "data", Data(sol.t, sol.u, v, c, params_les, params_dns))