Fix attention model

Author: SCiarella

configs/snellius64/att_base.yaml

@@ -0,0 +1,64 @@
+docreatedata: true
+docomp: true
+ntrajectory: 8
+T: "Float64"
+dataproj: true
+projtest: false
+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.RK44(; 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: "Attention"
+  type: attentioncnn
+  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]
+  use_attention: [true, false, false, false, false]
+  emb_sizes: [124, 124, 124, 124, 124]
+#  Ns: [148, 144, 140, 136, 132]
+  Ns: [ 84,  80,  76,  72,  68]
+#  patch_sizes: [37, 36, 35, 34, 33]
+  patch_sizes: [21, 20, 19, 18, 17]
+  n_heads: [4, 4, 4, 4, 4]
+  sum_attention: [false, false, false, false, false]
+  rng: "Xoshiro(seeds.θ_start)"
+priori:
+  dotrain: true
+  nepoch: 1000
+  batchsize: 64
+  opt: "Adam(T(1.0e-3))"
+  do_plot: false
+  plot_train: false
+  lambda: 0.00005
+posteriori:
+  dotrain: true
+  projectorders: "(ProjectOrder.Last, )"
+  nepoch: 3000
+  opt: "Adam(T(1.0e-4))"
+  nunroll: 5
+  nunroll_valid: 10
+  nsamples: 5
+  dt: 0.0001
+  do_plot: false
+  plot_train: false
+  lambda: 0.0000005
+  sciml_solver: "Tsit5()"

extra_model_workflow.jl

@@ -78,7 +78,13 @@ else
     logfile = "log_$(Dates.now()).out"
 end
 logfile = joinpath(logdir, logfile)
-setsnelliuslogger(logfile)
+# check if I am planning to use Enzyme, in which case I can not touch the logger
+if (haskey(conf["priori"], "ad_type") && occursin("Enzyme", conf["priori"]["ad_type"])) || 
+   (haskey(conf["posteriori"], "ad_type") && occursin("Enzyme", conf["posteriori"]["ad_type"]))
+    @warn "Enzyme is used, so logger will not be set to ConsoleLogger"
+else
+    setsnelliuslogger(logfile)
+end
 
 @info "# A-posteriori analysis: Forced turbulence (2D)"
 
@@ -98,6 +104,7 @@ using CairoMakie
 using CoupledNODE: loss_priori_lux, create_loss_post_lux
 using CUDA
 using DifferentialEquations
+using Enzyme
 using IncompressibleNavierStokes.RKMethods
 using JLD2
 using LaTeXStrings
@@ -111,6 +118,7 @@ using OptimizationOptimJL
 using OptimizationCMAEvolutionStrategy
 using ParameterSchedulers
 using Random
+using SciMLSensitivity
 
 
 # ## Random number seeds
@@ -174,6 +182,22 @@ dns_seeds_train = dns_seeds[1:ntrajectory-2]
 dns_seeds_valid = dns_seeds[ntrajectory-1:ntrajectory-1]
 dns_seeds_test = dns_seeds[ntrajectory:ntrajectory]
 
+doprojtest = conf["projtest"]
+if doprojtest && taskid == 1
+    testprojfile = joinpath(outdir, "test_dns_proj.jld2")
+    if isfile(testprojfile)
+        @info "Test DNS projection file already exists."
+    else
+        create_test_dns_proj(
+            nchunks = 8000;
+            params...,
+            rng = Xoshiro(2406),
+            backend = backend,
+            filename = testprojfile,
+        )
+    end
+end
+
 # Create data
 docreatedata = conf["docreatedata"]
 for i = 1:ntrajectory
@@ -247,7 +271,7 @@ let
     u = randn(T, params.nles[1], params.nles[1], 2, 10) |> device
     θ = θ_start |> device
     closure(u, θ, st)
-    gradient(θ -> sum(closure(u, θ, st)[1]), θ)
+    Zygote.gradient(θ -> sum(closure(u, θ, st)[1]), θ)
     clean()
 end
 
@@ -299,7 +323,8 @@ let
         plot_train = conf["priori"]["plot_train"],
         nepoch,
         dataproj = conf["dataproj"],
-        λ = conf["priori"]["lambda"],
+        λ = haskey(conf["priori"], "λ") ? eval(Meta.parse(conf["priori"]["λ"])) : nothing,
+        ad_type = haskey(conf["priori"], "ad_type") ? eval(Meta.parse(conf["priori"]["ad_type"])) : Optimization.AutoZygote(),
     )
 end
 end
@@ -412,7 +437,9 @@ let
         sensealg = sensealg,
         sciml_solver = sciml_solver,
         dataproj = conf["dataproj"],
-        λ = conf["posteriori"]["lambda"],
+        λ = haskey(conf["posteriori"], "λ") ? eval(Meta.parse(conf["posteriori"]["λ"])) : nothing,
+        multishoot_nt = haskey(conf["posteriori"], "multishoot_nt") ? conf["posteriori"]["multishoot_nt"] : 0,
+        ad_type = haskey(conf["posteriori"], "ad_type") ? eval(Meta.parse(conf["posteriori"]["ad_type"])) : Optimization.AutoZygote(),
     )
 end
 end

multisub.sh

@@ -22,3 +22,5 @@
 sbatch -J owr job_a100.sh configs/snellius64/cnn_owr.yaml
 sbatch -J bs3 job_a100.sh configs/snellius64/cnn_bs3.yaml
 #sbatch -J composite job_a100.sh configs/snellius64/cnn_composite.yaml
+
+sbatch -J att job_a100_extra.sh configs/snellius64/att_base.yaml