Merge pull request #126 from agdestein/tensorclosure

Tensor closure

lib/NeuralClosure/src/closure.jl

@@ -32,20 +32,6 @@ function create_closure(layers...; rng)
     closure, θ
 end
 
-"""
-Create tensor basis closure.
-"""
-function create_tensorclosure(layers...; setup, rng)
-    D = setup.grid.dimension()
-    cnn, θ = create_closure(layers...; rng)
-    function closure(u, θ)
-        B, V = tensorbasis(u, setup)
-        V = stack(V)
-        α = cnn(V, θ)
-        τ = sum(k -> α[ntuple(Returns(:), D)..., k] .* B[k], 1:length(B))
-    end
-end
-
 """
 Interpolate velocity components to volume centers.
 """

lib/NeuralClosure/src/training.jl

@@ -113,8 +113,7 @@ create_relerr_prior(f, x, y) = θ -> norm(f(x, θ) - y) / norm(y)
 """
 Create a-posteriori loss function.
 """
-function create_loss_post(; setup, method, psolver, closure, nsubstep = 1)
-    closure_model = wrappedclosure(closure, setup)
+function create_loss_post(; setup, method, psolver, closure_model, nsubstep = 1)
     setup = (; setup..., closure_model)
     (; Iu) = setup.grid
     inside = Iu[1]

lib/NeuralClosure/test/examplerun.jl

@@ -112,7 +112,7 @@
                 setup,
                 psolver,
                 method = RKMethods.RK44(),
-                m.closure,
+                closure_model = wrappedclosure(m.closure, setup),
                 nsubstep = 1,
             )
             dataloader = create_dataloader_post(

lib/PaperDC/src/train.jl

@@ -219,7 +219,7 @@ function trainpost(;
             setup,
             psolver,
             method = RKProject(params.method, projectorder),
-            closure,
+            closure_model = wrappedclosure(closure, setup),
             nsubstep, # Time steps per loss evaluation
         )
         data_train =

lib/SymmetryClosure/Project.toml

@@ -4,11 +4,13 @@ authors = ["Syver Døving Agdestein <[email protected]>"]
 version = "1.0.0"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -20,23 +22,25 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
 Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+ParameterSchedulers = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+WGLMakie = "276b4fcb-3e11-5398-bf8b-a0c2d153d008"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
-[sources.IncompressibleNavierStokes]
-path = "../.."
-
-[sources.NeuralClosure]
-path = "../NeuralClosure"
+[sources]
+IncompressibleNavierStokes = {path = "../.."}
+NeuralClosure = {path = "../NeuralClosure"}
 
 [compat]
+Accessors = "0.1"
 Adapt = "4"
 CUDA = "5"
 CairoMakie = "0.12"
-ChainRulesCore = "1.25.0"
+ChainRulesCore = "1"
 ComponentArrays = "0.15"
+Dates = "1"
 FFTW = "1"
 IncompressibleNavierStokes = "2"
 JLD2 = "0.5"
@@ -48,6 +52,8 @@ NNlib = "0.9"
 NeuralClosure = "1"
 Observables = "0.5"
 Optimisers = "0.3, 0.4"
-StaticArrays = "1.9.8"
+ParameterSchedulers = "0.4"
+StaticArrays = "1"
+WGLMakie = "0.10"
 Zygote = "0.6, 0.7"
 julia = "1.9"

lib/SymmetryClosure/generate_data.jl

@@ -1,4 +1,4 @@
-# LSP hack                                             #src
+# LSP hack (to get "go to definition" etc. working)    #src
 if false                                               #src
     include("src/SymmetryClosure.jl")                  #src
     include("../NeuralClosure/src/NeuralClosure.jl")   #src
@@ -18,86 +18,47 @@ end                                                    #src
 
 ########################################################################## #src
 
-@info "Loading packages"
-flush(stdout)
+# ## Setup
 
-using Adapt
-using CUDA
-using Dates
-using JLD2
-using NeuralClosure
-using Random
 using SymmetryClosure
+using Random
+using NeuralClosure
+using JLD2
 
-########################################################################## #src
-
-# SLURM specific variables
-jobid = haskey(ENV, "SLURM_JOB_ID") ? parse(Int, ENV["SLURM_JOB_ID"]) : nothing
-taskid =
-    haskey(ENV, "SLURM_ARRAY_TASK_ID") ? parse(Int, ENV["SLURM_ARRAY_TASK_ID"]) : nothing
-
-isnothing(jobid) || @info "Running on SLURM (jobid = $jobid)"
-isnothing(taskid) || @info "Task id = $taskid)"
+# Get SLURM specific variables (if any)
+(; jobid, taskid) = slurm_vars()
 
-########################################################################## #src
+# Log info about current run
+time_info()
 
-@info "Starting at $(Dates.now())"
-@info """
-Last commit:
+# Hardware selection
+(; backend, device, clean) = hardware()
 
-$(cd(() -> read(`git log -n 1`, String), @__DIR__))
-"""
+# Test case
+(; params, outdir, plotdir, seed_dns, ntrajectory) = testcase(backend)
 
-########################################################################## #src
-
-# ## Hardware selection
-
-if CUDA.functional()
-    ## For running on a CUDA compatible GPU
-    @info "Running on CUDA"
-    backend = CUDABackend()
-    CUDA.allowscalar(false)
-    device = x -> adapt(CuArray, x)
-    clean() = (GC.gc(); CUDA.reclaim())
-else
-    ## For running on CPU.
-    ## Consider reducing the sizes of DNS, LES, and CNN layers if
-    ## you want to test run on a laptop.
-    @warn "Running on CPU"
-    backend = CPU()
-    device = identity
-    clean() = nothing
-end
+# DNS seeds
+dns_seeds = splitseed(seed_dns, ntrajectory)
 
 ########################################################################## #src
 
-# ## Data generation
-#
-# Create filtered DNS data for training, validation, and testing.
-
-# Parameters
-case = SymmetryClosure.testcase()
+# ## Create data
 
-# DNS seeds
-ntrajectory = 8
-seeds = splitseed(case.seed_dns, ntrajectory)
-
-# Create data
-for (iseed, seed) in enumerate(seeds)
+for (iseed, seed) in enumerate(dns_seeds)
     if isnothing(taskid) || iseed == taskid
         @info "Creating DNS trajectory for seed $(repr(seed))"
     else
         # Each task does one initial condition
         @info "Skipping seed $(repr(seed)) for task $taskid"
         continue
     end
-    filenames = map(Iterators.product(params.nles, params.filters)) do nles, Φ
+    filenames = map(Iterators.product(params.nles, params.filters)) do (nles, Φ)
         f = getdatafile(outdir, nles, Φ, seed)
         datadir = dirname(f)
         ispath(datadir) || mkpath(datadir)
         f
     end
-    data = create_les_data(; case.params..., rng = Xoshiro(seed), filenames)
+    data = create_les_data(; params..., backend, rng = Xoshiro(seed), filenames)
     @info(
         "Trajectory info:",
         data[1].comptime / 60,
@@ -106,8 +67,11 @@ for (iseed, seed) in enumerate(seeds)
     )
 end
 
+########################################################################## #src
+
 # Computational time
-docomp = false
+
+docomp = true
 docomp && let
     comptime, datasize = 0.0, 0.0
     for seed in dns_seeds

lib/SymmetryClosure/src/SymmetryClosure.jl

@@ -1,14 +1,36 @@
 module SymmetryClosure
 
-using StaticArrays
+using Accessors
+using Adapt
 using ChainRulesCore
+using CUDA
+using Dates
+using IncompressibleNavierStokes
+using JLD2
 using KernelAbstractions
 using LinearAlgebra
-using IncompressibleNavierStokes
+using Lux
 using NeuralClosure
+using NNlib
+using Optimisers
+using Random
+using StaticArrays
 
 include("tensorclosure.jl")
+include("setup.jl")
+include("cases.jl")
+include("train.jl")
 
-export tensorclosure, polynomial
+export tensorclosure, polynomial, create_cnn
+export slurm_vars,
+    time_info,
+    hardware,
+    splatfileparts,
+    getdatafile,
+    namedtupleload,
+    splitseed,
+    getsetup,
+    testcase
+export trainpost
 
 end

lib/SymmetryClosure/src/cases.jl

@@ -1,25 +1,27 @@
-function testcase()
-    # Choose where to put output
+function testcase(backend)
     basedir = haskey(ENV, "DEEPDIP") ? ENV["DEEPDIP"] : joinpath(@__DIR__, "..")
-    outdir = mkpath(joinpath(basedir, "output", "kolmogorov"))
-
+    outdir = mkpath(joinpath(basedir, "output", "Kolmogorov2D"))
+    plotdir = mkpath(joinpath(outdir, "plots"))
     seed_dns = 123
+    ntrajectory = 8
     T = Float32
-
     params = (;
         D = 2,
         lims = (T(0), T(1)),
         Re = T(6e3),
         tburn = T(0.5),
-        tsim = T(5),
+        tsim = T(3),
         savefreq = 50,
-        ndns = 4096,
+        ndns = 2048,
         nles = [32, 64, 128],
         filters = [FaceAverage()],
+        backend,
         icfunc = (setup, psolver, rng) ->
             random_field(setup, T(0); kp = 20, psolver, rng),
-        method = RKMethods.LMWray3(; T),
+        method = LMWray3(; T),
         bodyforce = (dim, x, y, t) -> (dim == 1) * 5 * sinpi(8 * y),
         issteadybodyforce = true,
+        processors = (; log = timelogger(; nupdate = 100)),
     )
+    (; outdir, plotdir, seed_dns, ntrajectory, params)
 end

lib/SymmetryClosure/src/setup.jl

@@ -0,0 +1,61 @@
+# Some script utils
+
+function slurm_vars()
+    jobid = haskey(ENV, "SLURM_JOB_ID") ? parse(Int, ENV["SLURM_JOB_ID"]) : nothing
+    taskid =
+        haskey(ENV, "SLURM_ARRAY_TASK_ID") ? parse(Int, ENV["SLURM_ARRAY_TASK_ID"]) :
+        nothing
+    isnothing(jobid) || @info "Running on SLURM" jobid taskid
+    (; jobid, taskid)
+end
+
+function time_info()
+    @info "Starting at $(Dates.now())"
+    @info """
+    Last commit:
+
+    $(cd(() -> read(`git log -n 1`, String), @__DIR__))
+    """
+end
+
+hardware() =
+    if CUDA.functional()
+        @info "Running on CUDA"
+        CUDA.allowscalar(false)
+        backend = CUDABackend()
+        device = x -> adapt(backend, x)
+        clean = () -> (GC.gc(); CUDA.reclaim())
+        (; backend, device, clean)
+    else
+        @warn """
+        Running on CPU.
+        Consider reducing the size of DNS, LES, and CNN layers if
+        you want to test run on a laptop.
+        """
+        (; backend = CPU(), device = identity, clean = () -> nothing)
+    end
+
+function splatfileparts(args...; kwargs...)
+    sargs = string.(args)
+    skwargs = map((k, v) -> string(k) * "=" * string(v), keys(kwargs), values(kwargs))
+    s = [sargs..., skwargs...]
+    join(s, "_")
+end
+
+getdatafile(outdir, nles, filter, seed) =
+    joinpath(outdir, "data", splatfileparts(; seed = repr(seed), filter, nles) * ".jld2")
+
+function namedtupleload(file)
+    dict = load(file)
+    k, v = keys(dict), values(dict)
+    pairs = @. Symbol(k) => v
+    (; pairs...)
+end
+
+getsetup(; params, nles) = Setup(;
+    x = ntuple(α -> range(params.lims..., nles + 1), params.D),
+    params.Re,
+    params.backend,
+    params.bodyforce,
+    params.issteadybodyforce,
+)