Use adapt=false and RK4

Author: SCiarella

cnn_model_workflow.jl

@@ -349,6 +349,7 @@ let
         dns_seeds_train,
         dns_seeds_valid,
         nunroll = conf["posteriori"]["nunroll"],
+        dt = conf["posteriori"]["dt"],
         closure,
         closure_name,
         θ_start = θ_cnn_prior,
@@ -492,19 +493,20 @@ let
             t = sample.t[it],
         )
         tspan = (data.t[1], data.t[end])
+        dt = conf["posteriori"]["dt"]
 
         ## No model
         dudt_nomod = NS.create_right_hand_side_inplace(
             setup, psolver)
 
-        epost.nomodel[I,:], _ = compute_epost(dudt_nomod, θ_cnn_post[I].*0 , tspan, data, tsave)
+        epost.nomodel[I,:], _ = compute_epost(dudt_nomod, θ_cnn_post[I].*0 , tspan, data, tsave, dt)
         @info "Epost nomodel" epost.nomodel[I,:]
         # with closure
         dudt = NS.create_right_hand_side_with_closure_inplace(
             setup, psolver, closure, st)
-        epost.model_prior[I, :], _ = compute_epost(dudt, device(θ_cnn_prior[ig, ifil]) , tspan, data, tsave)
+        epost.model_prior[I, :], _ = compute_epost(dudt, device(θ_cnn_prior[ig, ifil]) , tspan, data, tsave, dt)
         @info "Epost model_prior" epost.model_prior[I, :]
-        epost.model_post[I, :], epost.model_t_post_inference[I] = compute_epost(dudt, device(θ_cnn_post[I]) , tspan, data, tsave)
+        epost.model_post[I, :], epost.model_t_post_inference[I] = compute_epost(dudt, device(θ_cnn_post[I]) , tspan, data, tsave, dt)
         @info "Epost model_post" epost.model_post[I, :]
 
         clean()

configs/local/cnn_1 copy.yaml

@@ -0,0 +1,50 @@
+docreatedata: true
+docomp: false
+ntrajectory: 8
+T: "Float32"
+params:
+  D: 2
+  lims: [0.0, 1.0]
+  Re: 6000.0
+  tburn: 0.5
+  tsim: 2.0
+  savefreq: 100
+  ndns: 1024
+  nles: [32]
+  filters: ["FaceAverage()"]
+  icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
+  method: "RKMethods.Wray3(; T)"
+  bodyforce: "(dim, x, y, t) -> (dim == 1) * 5 * sinpi(8 * y)"
+  issteadybodyforce: true
+  processors: "(; log = timelogger(; nupdate=100))"
+  Δt: 0.0001
+seeds:
+  dns: 1111
+  θ_start: 234
+  prior: 345
+  post: 456
+closure:
+  name: "cnn_new"
+  type: cnn
+  radii: [2, 2, 2, 2, 2]
+  channels: [24, 24, 24, 24, 2]
+  activations: ["tanh", "tanh", "tanh", "tanh", "identity"]
+  use_bias: [true, true, true, true, false]
+  rng: "Xoshiro(seeds.θ_start)"
+priori:
+  dotrain: true
+  nepoch: 5000
+  batchsize: 64
+  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  do_plot: false
+  plot_train: false
+posteriori:
+  dotrain: true
+  projectorders: "(ProjectOrder.Last, )"
+  nepoch: 300
+  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  nunroll: 5
+  nunroll_valid: 5
+  dt: T(1e-3)
+  do_plot: false
+  plot_train: false

configs/local/cnn_2.yaml

@@ -7,9 +7,9 @@ params:
   lims: [0.0, 1.0]
   Re: 6000.0
   tburn: 0.5
-  tsim: 2.0
+  tsim: 1.0
   savefreq: 100
-  ndns: 1024
+  ndns: 256
   nles: [32]
   filters: ["FaceAverage()"]
   icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
@@ -19,7 +19,7 @@ params:
   processors: "(; log = timelogger(; nupdate=100))"
   Δt: 0.0001
 seeds:
-  dns: 12345
+  dns: 1234567
   θ_start: 234
   prior: 345
   post: 456

extra_model_workflow.jl

@@ -368,6 +368,7 @@ let
         dns_seeds_train,
         dns_seeds_valid,
         nunroll = conf["posteriori"]["nunroll"],
+        dt = conf["posteriori"]["dt"],
         closure,
         closure_name,
         θ_start = θ_cnn_prior,
@@ -511,20 +512,20 @@ let
             u = selectdim(sample.u, ndims(sample.u), it) |> collect |> device,
             t = sample.t[it],
         )
-        dt = T(data.t[2] - data.t[1])
         tspan = (data.t[1], data.t[end])
+        dt = conf["posteriori"]["dt"]
 
         ## No model
         dudt_nomod = NS.create_right_hand_side_inplace(
             setup, psolver)
-        epost.nomodel[I, :], epost.nomodel_t_post_inference[I] = compute_epost(dudt_nomod, θ_cnn_post[I].*0 , tspan, data, tsave)
+        epost.nomodel[I, :], epost.nomodel_t_post_inference[I] = compute_epost(dudt_nomod, θ_cnn_post[I].*0 , tspan, data, tsave, dt)
         @info "Epost nomodel" epost.nomodel[I,:]
         # with closure
         dudt = NS.create_right_hand_side_with_closure_inplace(
             setup, psolver, closure, st)
-        epost.model_prior[I, :], _ = compute_epost(dudt, device(θ_cnn_prior[ig, ifil]) , tspan, data, tsave)
+        epost.model_prior[I, :], _ = compute_epost(dudt, device(θ_cnn_prior[ig, ifil]) , tspan, data, tsave, dt)
         @info "Epost model_prior" epost.model_prior[I, :]
-        epost.model_post[I, :], epost.model_t_post_inference[I] = compute_epost(dudt, device(θ_cnn_post[I]) , tspan, data, tsave)
+        epost.model_post[I, :], epost.model_t_post_inference[I] = compute_epost(dudt, device(θ_cnn_post[I]) , tspan, data, tsave, dt)
         @info "Epost model_post" epost.model_post[I, :]
         clean()
     end
@@ -657,7 +658,7 @@ let
         θ_prior = device(θ_cnn_prior[ig, ifil])
         θ_post = device(θ_cnn_post[I])
 
-        dt = T(sample.t[2] - sample.t[1])
+        dt = conf["posteriori"]["dt"]
         tspan = (sample.t[1], sample.t[end])
         dt_sample = T(0.05) # Sample every 0.05 seconds for the history (same as INS)
         tsave = (x*dt_sample for x in 1:(floor(Int, length(sample.t) / 0.05)+1))
@@ -671,10 +672,10 @@ let
         pred_prior =
                 solve(
                     prob_prior,
-                    Tsit5();
+                    RK4();
                     u0 = x,
                     p = θ_prior,
-                    adaptive = true,
+                    adaptive = false,
                     saveat = tsave,
                     dt = dt,
                     tspan = tspan,
@@ -685,10 +686,10 @@ let
         pred_post =
                 solve(
                     prob_post,
-                    Tsit5();
+                    RK4();
                     u0 = x,
                     p = θ_post,
-                    adaptive = true,
+                    adaptive = false,
                     saveat = tsave,
                     dt = dt,
                     tspan = tspan,
@@ -964,7 +965,7 @@ let
         θ_prior = device(θ_cnn_prior[I])
         θ_post = device(θ_cnn_post[I])
 
-        dt = T(1e-4)
+        dt = conf["posteriori"]["dt"]
         tspan = (T(0), times[end]+T(1e-4))
 
         dudt = NS.create_right_hand_side_with_closure_inplace(
@@ -976,23 +977,25 @@ let
         pred_prior =
             solve(
                     prob_prior,
-                    Tsit5(),
+                    RK4(),
                     u0 = x,
                     p = θ_prior,
-                    adaptive = true,
+                    adaptive = false,
                     saveat = times,
                     tspan = tspan,
+                    dt = dt,
             )
         prob_post = ODEProblem(dudt, x, tspan, θ_post)
         pred_post =
             solve(
                     prob_post,
-                    Tsit5(),
+                    RK4(),
                     u0 = x,
                     p = θ_post,
-                    adaptive = true,
+                    adaptive = false,
                     saveat = times,
                     tspan = tspan,
+                    dt = dt,
             )
 
         for it in 1:length(times)

src/train.jl

@@ -293,8 +293,9 @@ function trainpost(;
             dudt_nn,
             griddims,
             inside;
+            dt = dt,
             ensemble = nsamples > 1,
-            sciml_solver = Tsit5(),
+            sciml_solver = RK4(),
             sensealg = sensealg,
         )
 
@@ -386,7 +387,7 @@ function compute_t_prior_inference(closure, θ, st, x, y, nreps = 1000)
 end
 
 
-function compute_epost(rhs, ps, tspan, (u, t), tsave)
+function compute_epost(rhs, ps, tspan, (u, t), tsave, dt)
     griddims = ((:) for _ = 1:(ndims(u)-2))
     inside = ((2:(size(u, 1)-1)) for _ = 1:(ndims(u)-2))
     x = u[griddims..., :, 1]
@@ -395,13 +396,14 @@ function compute_epost(rhs, ps, tspan, (u, t), tsave)
     t0 = time()
     pred = solve(
         prob,
-        Tsit5();
+        RK4();
         u0 = x,
         p = ps,
-        adaptive = true,
+        adaptive = false,
         saveat = Array(t),
         tspan = tspan,
         save_start = false,
+        dt = dt,
     )
 
     e = 0.0