From 03acafa3566742b1f37d722d2207f5dc21d031c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Syver=20D=C3=B8ving=20Agdestein?= Date: Mon, 9 Dec 2024 15:56:21 +0100 Subject: [PATCH] fix: use elementwise operations --- lib/SymmetryClosure/generate_data.jl | 126 +++++++++++++++++++++++ lib/SymmetryClosure/src/cases.jl | 25 +++++ lib/SymmetryClosure/src/tensorclosure.jl | 37 ++++--- lib/SymmetryClosure/test_tensor.jl | 50 +++++++-- 4 files changed, 216 insertions(+), 22 deletions(-) create mode 100644 lib/SymmetryClosure/generate_data.jl create mode 100644 lib/SymmetryClosure/src/cases.jl diff --git a/lib/SymmetryClosure/generate_data.jl b/lib/SymmetryClosure/generate_data.jl new file mode 100644 index 00000000..ac2c73df --- /dev/null +++ b/lib/SymmetryClosure/generate_data.jl @@ -0,0 +1,126 @@ +# LSP hack #src +if false #src + include("src/SymmetryClosure.jl") #src + include("../NeuralClosure/src/NeuralClosure.jl") #src + include("../../src/IncompressibleNavierStokes.jl") #src + using .SymmetryClosure #src + using .NeuralClosure #src + using .IncompressibleNavierStokes #src +end #src + +########################################################################## #src + +# # Data generation +# +# This script is used to generate filtered DNS data. + +@info "# DNS data generation" + +########################################################################## #src + +@info "Loading packages" +flush(stdout) + +using Adapt +using CUDA +using Dates +using JLD2 +using NeuralClosure +using Random +using SymmetryClosure + +########################################################################## #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)" + +########################################################################## #src + +@info "Starting at $(Dates.now())" +@info """ +Last commit: + +$(cd(() -> read(`git log -n 1`, String), @__DIR__)) +""" + +########################################################################## #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 + +########################################################################## #src + +# ## Data generation +# +# Create filtered DNS data for training, validation, and testing. + +# Parameters +case = SymmetryClosure.testcase() + +# DNS seeds +ntrajectory = 8 +seeds = splitseed(case.seed_dns, ntrajectory) + +# Create data +for (iseed, seed) in enumerate(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, Φ + f = getdatafile(outdir, nles, Φ, seed) + datadir = dirname(f) + ispath(datadir) || mkpath(datadir) + f + end + data = create_les_data(; case.params..., rng = Xoshiro(seed), filenames) + @info( + "Trajectory info:", + data[1].comptime / 60, + length(data[1].t), + Base.summarysize(data) * 1e-9, + ) +end + +# Computational time +docomp = false +docomp && let + comptime, datasize = 0.0, 0.0 + for seed in dns_seeds + comptime += load( + getdatafile(outdir, params.nles[1], params.filters[1], seed), + "comptime", + ) + end + for seed in dns_seeds, nles in params.nles, Φ in params.filters + data = namedtupleload(getdatafile(outdir, nles, Φ, seed)) + datasize += Base.summarysize(data) + end + @info "Data" comptime + @info "Data" comptime / 60 datasize * 1e-9 + clean() +end diff --git a/lib/SymmetryClosure/src/cases.jl b/lib/SymmetryClosure/src/cases.jl new file mode 100644 index 00000000..f2f31a05 --- /dev/null +++ b/lib/SymmetryClosure/src/cases.jl @@ -0,0 +1,25 @@ +function testcase() + # Choose where to put output + basedir = haskey(ENV, "DEEPDIP") ? ENV["DEEPDIP"] : joinpath(@__DIR__, "..") + outdir = mkpath(joinpath(basedir, "output", "kolmogorov")) + + seed_dns = 123 + T = Float32 + + params = (; + D = 2, + lims = (T(0), T(1)), + Re = T(6e3), + tburn = T(0.5), + tsim = T(5), + savefreq = 50, + ndns = 4096, + nles = [32, 64, 128], + filters = [FaceAverage()], + icfunc = (setup, psolver, rng) -> + random_field(setup, T(0); kp = 20, psolver, rng), + method = RKMethods.LMWray3(; T), + bodyforce = (dim, x, y, t) -> (dim == 1) * 5 * sinpi(8 * y), + issteadybodyforce = true, + ) +end diff --git a/lib/SymmetryClosure/src/tensorclosure.jl b/lib/SymmetryClosure/src/tensorclosure.jl index 1dfb55e2..23c8bc1e 100644 --- a/lib/SymmetryClosure/src/tensorclosure.jl +++ b/lib/SymmetryClosure/src/tensorclosure.jl @@ -1,8 +1,13 @@ function tensorclosure(model, u, θ, setup) + (; Δ) = setup.grid + @assert all(Δi -> all(≈(Δi[1]), Δi), Array.(Δ)) + Δx = Array(Δ[1])[1] B, V = tensorbasis(u, setup) x = model(V, θ) + x .*= Δx^2 # τ = sum(x .* B; dims = ndims(B)) τ = IncompressibleNavierStokes.lastdimcontract(x, B, setup) + τ = apply_bc_p(τ, zero(eltype(u)), setup) IncompressibleNavierStokes.divoftensor(τ, setup) end @@ -10,7 +15,7 @@ function polynomial(V, θ) s..., nV = size(V) V = eachslice(V; dims = ndims(V)) basis = if nV == 2 - cat(V[1], V[2], V[1] * V[2], V[1]^2, V[2]^2; dims = length(s) + 1) + cat(V[1], V[2], V[1] .* V[2], V[1] .^ 2, V[2] .^ 2; dims = length(s) + 1) elseif nV == 5 cat( V[1], @@ -18,21 +23,21 @@ function polynomial(V, θ) V[3], V[4], V[5], - V[1] * V[2], - V[1] * V[3], - V[1] * V[4], - V[1] * V[5], - V[2] * V[3], - V[2] * V[4], - V[2] * V[5], - V[3] * V[4], - V[3] * V[5], - V[4] * V[5], - V[1]^2, - V[2]^2, - V[3]^2, - V[4]^2, - V[5]^2; + V[1] .* V[2], + V[1] .* V[3], + V[1] .* V[4], + V[1] .* V[5], + V[2] .* V[3], + V[2] .* V[4], + V[2] .* V[5], + V[3] .* V[4], + V[3] .* V[5], + V[4] .* V[5], + V[1] .^ 2, + V[2] .^ 2, + V[3] .^ 2, + V[4] .^ 2, + V[5] .^ 2; dims = length(s) + 1, ) end diff --git a/lib/SymmetryClosure/test_tensor.jl b/lib/SymmetryClosure/test_tensor.jl index 69ec7193..edd39ea0 100644 --- a/lib/SymmetryClosure/test_tensor.jl +++ b/lib/SymmetryClosure/test_tensor.jl @@ -3,6 +3,7 @@ if false using .SymmetryClosure end +# using CairoMakie using CairoMakie using IncompressibleNavierStokes using SymmetryClosure @@ -10,21 +11,23 @@ using CUDA using Zygote using Random using LinearAlgebra +using NeuralClosure lines(cumsum(randn(100))) # Setup -n = 8 -# ax = range(0.0, 1.0, n + 1) -# x = ax, ax -x = tanh_grid(0.0, 1.0, n + 1), stretched_grid(-0.2, 1.2, n + 1) +n = 128 +ax = range(0.0, 1.0, n + 1) +x = ax, ax +# x = tanh_grid(0.0, 1.0, n + 1), stretched_grid(-0.2, 1.2, n + 1) setup = Setup(; x, Re = 1e4, backend = CUDABackend(), - boundary_conditions = ((DirichletBC(), DirichletBC()), (DirichletBC(), DirichletBC())), + # boundary_conditions = ((DirichletBC(), DirichletBC()), (DirichletBC(), DirichletBC())), ); -ustart = vectorfield(setup) |> randn! +ustart = random_field(setup) +# ustart = vectorfield(setup) |> randn! u = ustart @@ -65,7 +68,42 @@ let nothing end +urot = NeuralClosure.rot2stag(u, 1) B, V = tensorbasis(u, setup) +Brot, Vrot = tensorbasis(urot, setup) + +iV = 2 +# V[:, :, iV] |> Array |> heatmap +Vrot[:, :, iV] |> Array |> heatmap +rot2(V, 1)[:, :, iV] |> Array |> heatmap + +Vrot - rot2(V, 1) .|> abs |> maximum + +begin + θ = randn(5, 3) + # θ[:, 1] .= 0 + # θ[:, 3] .= 0 + # θ = zeros(5, 3) + # θ[3, :] .= 1e-3 * randn() + θ = θ |> CuArray + s = tensorclosure(polynomial, u, θ, setup) |> u -> apply_bc_u(u, zero(eltype(u)), setup) + srot = + tensorclosure(polynomial, urot, θ, setup) |> + u -> apply_bc_u(u, zero(eltype(u)), setup) + rots = NeuralClosure.rot2stag(s, 1) + (srot-rots)[2:end-1, 2:end-1, :] .|> abs |> maximum +end + +i = 1 +srot[2:end-1, 2:end-1, i] |> Array |> heatmap +rots[2:end-1, 2:end-1, i] |> Array |> heatmap +# s[2:end-1, 2:end-1, 2] |> Array |> heatmap +(srot-rots)[2:end-1, 2:end-1, i] |> Array |> heatmap +(srot-rots)[2:end-1, 2:end-1, :] .|> abs |> maximum +he = (srot-rots)[2:end-1, 2:end-1, :] +he[125:128, 1:10, 1] + +x = randn(5, 5, 2) typeof(B) getindex.(B, 1)