From 7665b1138c9fa24a4a5a3e410181fbc41a8da9b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Syver=20D=C3=B8ving=20Agdestein?= <syverda@gmail.com>
Date: Thu, 21 Dec 2023 09:22:10 +0100
Subject: [PATCH] Add different setup

---
 scratch/multigrid.jl            |  5 +++--
 scratch/train_model.jl          |  7 ++++---
 src/closures/create_les_data.jl | 33 ++++++++++++++++++++++-----------
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/scratch/multigrid.jl b/scratch/multigrid.jl
index e3f4e7061..1e470d10e 100644
--- a/scratch/multigrid.jl
+++ b/scratch/multigrid.jl
@@ -59,9 +59,10 @@ get_params(nles) = (;
     tburn = T(0.05),
     tsim = T(0.5),
     Δt = T(1e-4),
-    nles,
-    ndns = 2048,
+    nles = map(n -> (n, n), nles),
+    ndns = (2048, 2048),
     ArrayType,
+    PSolver = SpectralPressureSolver,
 )
 
 params_train = (; get_params([64, 128, 256])..., savefreq = 5);
diff --git a/scratch/train_model.jl b/scratch/train_model.jl
index 916dd982c..a03c1734f 100644
--- a/scratch/train_model.jl
+++ b/scratch/train_model.jl
@@ -57,14 +57,15 @@ params = (;
     D = 2,
     Re = T(6_000),
     lims = (T(0), T(1)),
-    nles = [nles],
-    # ndns = 512,
-    ndns = 1024,
+    nles = [(nles, nles)],
+    # ndns = (512, 512),
+    ndns = (1024, 1024),
     tburn = T(0.05),
     tsim = T(0.5),
     Δt = T(1e-4),
     savefreq = 5,
     ArrayType,
+    PSolver = SpectralPressureSolver,
 )
 
 # Create LES data from DNS
diff --git a/src/closures/create_les_data.jl b/src/closures/create_les_data.jl
index 67493ba60..9e75b2674 100644
--- a/src/closures/create_les_data.jl
+++ b/src/closures/create_les_data.jl
@@ -105,30 +105,41 @@ function create_les_data(
     T;
     D = 2,
     Re = T(2_000),
-    lims = (T(0), T(1)),
-    nles = [64],
-    ndns = 256,
+    lims = ntuple(α -> (T(0), T(1)), D),
+    nles = [ntuple(α -> 64, D)],
+    ndns = ntuple(α -> 256, D),
     tburn = T(0.1),
     tsim = T(0.1),
     Δt = T(1e-4),
+    PSolver = SpectralPressureSolver,
     savefreq = 1,
     ArrayType = Array,
-    ic_params = (;),
+    icfunc = (setup, psolver) -> random_field(setup, T(0); psolver),
+    boundary_conditions = ntuple(α -> (PeriodicBC(), PeriodicBC()), D),
 )
     compression = @. ndns ÷ nles
     @assert all(@.(compression * nles == ndns))
 
     # Build setup and assemble operators
-    dns = Setup(ntuple(α -> LinRange(lims..., ndns + 1), D)...; Re, ArrayType)
+    dns = Setup(
+        ntuple(α -> LinRange(lims[α]..., ndns[α] + 1), D)...;
+        Re,
+        boundary_conditions,
+        ArrayType,
+    )
     les = [
-        Setup(ntuple(α -> LinRange(lims..., nles + 1), D)...; Re, ArrayType) for
-        nles in nles
+        Setup(
+            ntuple(α -> LinRange(lims[α], nles[α] + 1), D)...;
+            Re,
+            boundary_conditions,
+            ArrayType,
+        ) for nles in nles
     ]
 
     # Since the grid is uniform and identical for x and y, we may use a specialized
     # spectral pressure solver
-    psolver = SpectralPressureSolver(dns)
-    psolver_les = SpectralPressureSolver.(les)
+    psolver = PSolver(dns)
+    psolver_les = PSolver.(les)
 
     # Number of time steps to save
     nt = round(Int, tsim / Δt)
@@ -136,11 +147,11 @@ function create_les_data(
 
     # datasize = Base.summarysize(filtered) / 1e6
     datasize =
-        (nt ÷ savefreq + 1) * sum(nles .^ D) * 3 * 2 * length(bitstring(zero(T))) / 8 / 1e6
+        (nt ÷ savefreq + 1) * sum(prod.(nles)) * D * 2 * length(bitstring(zero(T))) / 8 / 1e6
     @info "Generating $datasize Mb of LES data"
 
     # Initial conditions
-    u₀ = random_field(dns, T(0); psolver, ic_params...)
+    u₀ = icfunc(dns, psolver)
 
     # Random body force
     # force_dns =