Add dataproj key to use same ref seed

Author: SCiarella

cnn_model_workflow.jl

@@ -169,7 +169,7 @@ dns_seeds_test = dns_seeds[ntrajectory:ntrajectory]
 docreatedata = conf["docreatedata"]
 for i = 1:ntrajectory
 	if i%numtasks == taskid - 1
-		docreatedata && createdata(; params, seed = dns_seeds[i], outdir, backend)
+		docreatedata && createdata(; params, seed = dns_seeds[i], outdir, backend, dataproj = conf["dataproj"])
 	end
 end
 @info "Data generated"
@@ -271,7 +271,8 @@ let
         batchsize = conf["priori"]["batchsize"],
         do_plot = conf["priori"]["do_plot"],
         plot_train = conf["priori"]["plot_train"],
-        nepoch,
+        nepoch = nepoch,
+        dataproj = conf["dataproj"],
     )
 end
 end
@@ -369,6 +370,7 @@ let
         do_plot = conf["posteriori"]["do_plot"],
         plot_train = conf["posteriori"]["plot_train"],
         sensealg = haskey(conf["posteriori"],:sensealg) ? eval(Meta.parse(conf["posteriori"]["sensealg"])) : nothing,
+        dataproj = conf["dataproj"],
     )
 end
 end

configs/snellius/cnn_1.yaml

@@ -2,6 +2,7 @@ docreatedata: true
 docomp: true
 ntrajectory: 8
 T: "Float32"
+dataproj: true
 params:
   D: 2
   lims: [0.0, 1.0]

configs/snellius/cnn_ins.yaml

@@ -2,6 +2,7 @@ docreatedata: true
 docomp: true
 ntrajectory: 8
 T: "Float32"
+dataproj: false
 params:
   D: 2
   lims: [0.0, 1.0]
@@ -48,3 +49,4 @@ posteriori:
   dt: 0.0001
   do_plot: false
   plot_train: false
+  nsamples: 1

src/plots.jl

@@ -236,7 +236,7 @@ function plot_energy_evolution(
         end
     end
 
-    if closure_name == "cnn_1"
+    if closure_name == "cnn_proj"
         label = "No closure (projected dyn)"
         if _missing_label(ax, label) && haskey(energyhistory, Symbol("nomodel"))
             lines!(
@@ -702,7 +702,7 @@ function plot_epost_vs_t(error_file, closure_name, nles, ax, color, PLOT_STYLES)
         end
     end
 
-    if closure_name == "cnn_INS"
+    if closure_name == "cnn_noproj"
         label = "No model (projected dyn)"
         if _missing_label(ax, label)  # add No closure only once
             scatterlines!(

src/train.jl

@@ -3,19 +3,26 @@ function getdatafile(outdir, nles, filter, seed)
     joinpath(outdir, "data", splatfileparts(; seed = repr(seed), filter, nles) * ".jld2")
 end
 
-function load_data_set(outdir, nles, Φ, seeds)
+function load_data_set(outdir, nles, Φ, seeds, dataproj)
     data = []
     for s in seeds
-        data_i = namedtupleload(getdatafile(outdir, nles, Φ, s))
+        filename = getdatafile(outdir, nles, Φ, s)
+        if dataproj
+            filename = replace(filename, ".jld2" => "_projected.jld2")
+        end
+        data_i = namedtupleload(filename)
         push!(data, data_i)
     end
     return data
 end
 
-function createdata(; params, seed, outdir, backend)
+function createdata(; params, seed, outdir, backend, dataproj)
     for (nles, Φ) in Iterators.product(params.nles, params.filters)
 
         filename = getdatafile(outdir, nles, Φ, seed)
+        if dataproj
+            filename = replace(filename, ".jld2" => "_projected.jld2")
+        end
         datadir = dirname(filename)
         ispath(datadir) || mkpath(datadir)
 
@@ -85,6 +92,7 @@ function trainprior(;
     do_plot = false,
     plot_train = false,
     nepoch,
+    dataproj
 )
     device(x) = adapt(params.backend, x)
     itotal = 0
@@ -118,8 +126,8 @@ function trainprior(;
         NS = Base.get_extension(CoupledNODE, :NavierStokes)
 
         # Read the data in the format expected by the CoupledNODE
-        data_train = load_data_set(outdir, nles, Φ, dns_seeds_train)
-        data_valid = load_data_set(outdir, nles, Φ, dns_seeds_valid)
+        data_train = load_data_set(outdir, nles, Φ, dns_seeds_train, dataproj)
+        data_valid = load_data_set(outdir, nles, Φ, dns_seeds_valid, dataproj)
         @assert length(nles) == 1 "Only one nles for a-priori training"
         io_train = NS.create_io_arrays_priori(data_train, setup[1], device)
         io_valid = NS.create_io_arrays_priori(data_valid, setup[1], device)
@@ -244,6 +252,7 @@ function trainpost(;
     do_plot = false,
     plot_train = false,
     sensealg = nothing,
+    dataproj,
 )
     device(x) = adapt(params.backend, x)
     itotal = 0
@@ -280,8 +289,8 @@ function trainpost(;
         end
 
         # Read the data in the format expected by the CoupledNODE
-        data_train = load_data_set(outdir, nles, Φ, dns_seeds_train)
-        data_valid = load_data_set(outdir, nles, Φ, dns_seeds_valid)
+        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)