autopush

Author: SCiarella

Project.toml

@@ -43,7 +43,6 @@ GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 AttentionLayer = {rev = "main", url = "https://github.com/DEEPDIP-project/AttentionLayer.jl.git"}
 ConvolutionalNeuralOperators = {rev = "main", url = "https://github.com/DEEPDIP-project/ConvolutionalNeuralOperators.jl.git"}
 NeuralClosure = {rev = "main", url = "https://github.com/DEEPDIP-project/NeuralClosure.jl.git"}
-CoupledNODE = {rev = "main", url = "https://github.com/DEEPDIP-project/CoupledNODE.jl.git"}
 
 [compat]
 Accessors = "0.1"

cnn_model_workflow.jl

@@ -357,6 +357,7 @@ let
         postseed = seeds.post,
         dns_seeds_train,
         dns_seeds_valid,
+        dns_seeds_test,
         nunroll = conf["posteriori"]["nunroll"],
         nsamples = conf["posteriori"]["nsamples"],
         dt = T(conf["posteriori"]["dt"]),
@@ -497,7 +498,8 @@ let
         @info "Computing a-posteriori errors" projectorder Φ nles
         I = CartesianIndex(ig, ifil, iorder)
         setup = getsetup(; params, nles)
-        psolver = psolver_spectral(setup)
+        #psolver = psolver_spectral(setup)
+        psolver = default_psolver(setup)
         sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1]))
         it = 1:length(sample.t)
         data = (;
@@ -520,7 +522,7 @@ let
             setup, psolver, closure, st)
         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, dt)
+        epost.model_post[I, :], epost.model_t_post_inference[I] = compute_epost(dudt, device(θ_cnn_post[ig, ifil, iorder]) , tspan, data, tsave, dt)
         @info "Epost model_post" epost.model_post[I, :]
 
         clean()

configs/snellius/cnn_1.yaml

@@ -34,15 +34,15 @@ closure:
   rng: "Xoshiro(seeds.θ_start)"
 priori:
   dotrain: true
-  nepoch: 10000
+  nepoch: 50000
   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
+  nepoch: 1500
   opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
   nunroll: 5
   nunroll_valid: 10

configs/snellius/cnn_2.yaml

@@ -2,6 +2,7 @@ docreatedata: true
 docomp: true
 ntrajectory: 8
 T: "Float32"
+dataproj: true
 params:
   D: 2
   lims: [0.0, 1.0]
@@ -42,7 +43,7 @@ priori:
 posteriori:
   dotrain: true
   projectorders: "(ProjectOrder.Last, )"
-  nepoch: 300
+  nepoch: 100
   opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
   nunroll: 5
   nunroll_valid: 10

configs/snellius/cnn_ins.yaml

@@ -34,15 +34,15 @@ closure:
   rng: "Xoshiro(seeds.θ_start)"
 priori:
   dotrain: true
-  nepoch: 10000
+  nepoch: 50000
   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
+  nepoch: 1500
   opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
   nunroll: 5
   nunroll_valid: 10

configs/snellius/cnn_nt25.yaml

@@ -0,0 +1,53 @@
+docreatedata: true
+docomp: true
+ntrajectory: 8
+T: "Float32"
+dataproj: true
+params:
+  D: 2
+  lims: [0.0, 1.0]
+  Re: 6000.0
+  tburn: 0.5
+  tsim: 5.0
+  savefreq: 50
+  ndns: 4096
+  nles: [64]
+  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.00005
+seeds:
+  dns: 123456
+  θ_start: 234
+  prior: 345
+  post: 456
+closure:
+  name: "cnn_project_nt25"
+  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:
+  reuse: "cnn_project"
+  dotrain: true
+  nepoch: 50000
+  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: 1500
+  opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
+  nunroll: 25
+  nunroll_valid: 10
+  nsamples: 1
+  dt: 0.0001
+  do_plot: false
+  plot_train: false

configs/snellius/cnn_test.yaml

@@ -0,0 +1,53 @@
+docreatedata: true
+docomp: true
+ntrajectory: 8
+T: "Float32"
+dataproj: true
+params:
+  D: 2
+  lims: [0.0, 1.0]
+  Re: 6000.0
+  tburn: 0.5
+  tsim: 5.0
+  savefreq: 50
+  ndns: 4096
+  nles: [64]
+  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.00005
+seeds:
+  dns: 123456
+  θ_start: 234
+  prior: 345
+  post: 456
+closure:
+  name: "cnn_test"
+  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:
+  reuse: "cnn_project"
+  dotrain: true
+  nepoch: 50000
+  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: 100
+  opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.1))"
+  nunroll: 5
+  nunroll_valid: 24
+  nsamples: 5
+  dt: 0.00005
+  do_plot: false
+  plot_train: false

src/train.jl

@@ -239,6 +239,7 @@ function trainpost(;
     postseed,
     dns_seeds_train,
     dns_seeds_valid,
+    dns_seeds_test,
     nunroll,
     nsamples = 1,
     closure,
@@ -281,21 +282,15 @@ function trainpost(;
         checkfile = join(splitext(postfile), "_checkpoint")
         setup = getsetup(; params, nles)
         psolver = default_psolver(setup)
-        # Read the data in the format expected by the CoupledNODE
         T = eltype(params.Re)
-        setup = []
-        for nl in nles
-            x = ntuple(α -> LinRange(T(0.0), T(1.0), nl + 1), params.D)
-            push!(setup, Setup(; x = x, Re = params.Re, params.backend))
-        end
 
         # Read the data in the format expected by the CoupledNODE
         data_train = load_data_set(outdir, nles, Φ, dns_seeds_train, dataproj)
         data_valid = load_data_set(outdir, nles, Φ, dns_seeds_valid, dataproj)
 
         NS = Base.get_extension(CoupledNODE, :NavierStokes)
-        io_train = NS.create_io_arrays_posteriori(data_train, setup[1], device)
-        io_valid = NS.create_io_arrays_posteriori(data_valid, setup[1], device)
+        io_train = NS.create_io_arrays_posteriori(data_train, setup, device)
+        io_valid = NS.create_io_arrays_posteriori(data_valid, setup, device)
         θ = device(copy(θ_start[itotal]))
         dataloader_post = NS.create_dataloader_posteriori(
             io_train;
@@ -305,7 +300,7 @@ function trainpost(;
             device = device,
         )
 
-        dudt_nn = NS.create_right_hand_side_with_closure(setup[1], psolver, closure, st)
+        dudt_nn = NS.create_right_hand_side_with_closure(setup, psolver, closure, st)
         griddims = ((:) for _ = 1:params.D)
         inside = ((2:(nles+1)) for _ = 1:params.D)
         loss = CoupledNODE.create_loss_post_lux(
@@ -338,11 +333,25 @@ function trainpost(;
             nepochs_left = nepoch
         end
 
+
+        # For the callback I am going to use the a-posteriori error estimator
+        sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1]))
+        it = 1:(nunroll_valid+1)
+        data_cb = (;
+            u = selectdim(sample.u, ndims(sample.u), it) |> collect |> device,
+            t = sample.t[it],
+        )
+        tspan = (data_cb.t[1], data_cb.t[end])
+        tsave = [nunroll_valid]
+        dudt_cb = NS.create_right_hand_side_with_closure_inplace(
+            setup, psolver, closure, st)
+        loss_cb(_model, pp, _st, _data ) = compute_epost(dudt_cb, pp , tspan, data_cb, tsave, dt)[1][end]
+
         callbackstate, callback = NS.create_callback(
             closure,
             θ,
             io_valid,
-            loss,
+            loss_cb,
             st;
             callbackstate = callbackstate,
             nunroll = nunroll_valid,
@@ -420,9 +429,10 @@ function compute_epost(rhs, ps, tspan, (u, t), tsave, dt)
         p = ps,
         adaptive = true,
         saveat = Array(t),
-        tspan = tspan,
-        save_start = false,
-        dt = dt,
+        #tstops = Array(t),
+        #tspan = tspan,
+        #save_start = false,
+        #dt = dt,
     )
 
     e = 0.0
@@ -441,6 +451,12 @@ function compute_epost(rhs, ps, tspan, (u, t), tsave, dt)
             push!(es, e / (it - 1))
         end
     end
+    #for it in tsave
+    #    yref = y[inside..., :, 1:it]
+    #    ypred = pred[inside..., :, 1:it]
+
+    #    Lux.MSELoss()(ypred, yref) |> e -> push!(es, e)
+    #end
 
     return es, time() - t0