diff --git a/scratch/train_model.jl b/scratch/train_model.jl index a3d2edbdc..84e87fd80 100644 --- a/scratch/train_model.jl +++ b/scratch/train_model.jl @@ -68,8 +68,8 @@ size(data_train.u[1][1][1]) o = Observable(data_train.u[1][end][1]) heatmap(o) for i = 1:501 - # o[] = data_train.u[1][i][1] - o[] = data_train.cF[1][i][1] + o[] = data_train.u[1][i][1] + # o[] = data_train.cF[1][i][1] sleep(0.001) end @@ -98,7 +98,7 @@ closure, θ₀ = cnn( # Bias [true, true, true, false]; -) +); # closure, θ₀ = fno( # setup, @@ -116,28 +116,38 @@ closure, θ₀ = cnn( # gelu, # ); -@info "Closure model has $(length(θ₀)) parameters" +closure.NN + +# Create input/output arrays +function create_io_arrays(data, setup) + nsample = length(data.u) + nt = length(data.u[1]) - 1 + D = setup.grid.dimension() + T = eltype(data.u[1][1][1]) + (; N) = setup.grid + u = zeros(T, (N .- 2)..., D, nt + 1, nsample) + c = zeros(T, (N .- 2)..., D, nt + 1, nsample) + ifield = ntuple(Returns(:), D) + for i = 1:nsample, j = 1:nt+1, α = 1:D + copyto!(view(u, ifield..., α, j, i), view(data.u[i][j][α], setup.grid.Iu[α])) + copyto!(view(c, ifield..., α, j, i), view(data.cF[i][j][α], setup.grid.Iu[α])) + end + reshape(u, (N .- 2)..., D, :), reshape(c, (N .- 2)..., D, :) +end -# Test data -u_test = device(reshape(data_test.u[:, 1:20, 1:2], :, 40)) -c_test = device(reshape(data_test.cF[:, 1:20, 1:2], :, 40)) +io_train = create_io_arrays(data_train, setup) +io_valid = create_io_arrays(data_valid, setup) +io_test = create_io_arrays(data_test, setup) # Prepare training -θ = 1.0f-1 * device(θ₀) -# θ = device(θ₀) +θ = 1.0f-1 * cu(θ₀) +# θ = cu(θ₀) opt = Optimisers.setup(Adam(1.0f-3), θ) callbackstate = Point2f[] -randloss = create_randloss( - mean_squared_error, - closure, - data_train.V, - data_train.cF; - nuse = 50, - device, -) +randloss = create_randloss(mean_squared_error, closure, io_train...; nuse = 50, device = cu) # Warm-up -randloss(θ); +randloss(θ) @time randloss(θ); first(gradient(randloss, θ)); @time first(gradient(randloss, θ)); @@ -153,23 +163,21 @@ first(gradient(randloss, θ)); niter = 2000, ncallback = 10, callbackstate, - callback = create_callback(closure, u_test, c_test; state = callbackstate), -) + callback = create_callback(closure, cu(io_valid)...; state = callbackstate), +); GC.gc() CUDA.reclaim() Array(θ) # # Save trained parameters -# jldsave("output/forced/theta_cnn.jld2"; θ = Array(θ)) -# jldsave("output/forced/theta_fno.jld2"; θ = Array(θ)) +# jldsave("output/forced/theta_cnn.jld2"; theta = Array(θ)) +# jldsave("output/forced/theta_fno.jld2"; theta = Array(θ)) # # Load trained parameters # θθ = load("output/theta_cnn.jld2") # θθ = load("output/theta_fno.jld2") -# θθ = θθ["θ"] -# θθ = cu(θθ) -# θ .= θθ +# copyto!(θ, θθ["theta"]) relative_error(closure(device(data_train.V[:, 1, :]), θ), device(data_train.cF[:, 1, :])) relative_error( diff --git a/src/closures/create_les_data.jl b/src/closures/create_les_data.jl index c97b0d586..6e12f9261 100644 --- a/src/closures/create_les_data.jl +++ b/src/closures/create_les_data.jl @@ -67,7 +67,7 @@ _filter_saver(dns, les, comp; nupdate = 1) = processor(function (state) push!(_t, t) push!(_u, Array.(ubar)) # push!(_p, Array(pbar)) - push!(_F, Array.(Fubar)) + # push!(_F, Array.(Fubar)) # push!(_FG, Array.(FGbar)) push!(_cF, Array.(cF)) # push!(_cFG, Array.(cFG)) @@ -77,7 +77,7 @@ _filter_saver(dns, les, comp; nupdate = 1) = processor(function (state) t = _t, u = _u, # p = _p, - F = _F, + # F = _F, # FG = _FG, cF = _cF, # cFG = _cFG, @@ -119,17 +119,18 @@ function create_les_data( Δt, u = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), # p = fill(fill(zeros(T, N...), 0), 0), - F = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), + # F = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), # FG = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), cF = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), # cFG = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), # force = fill(fill(ntuple(α -> zeros(T, N...), D), 0), 0), ) - @info "Generating $(Base.summarysize(filtered) / 1e6) Mb of LES data" + # @info "Generating $(Base.summarysize(filtered) / 1e6) Mb of LES data" + @info "Generating $(nsim * (nt + 1) * nles * 3 * 2 * length(bitstring(zero(T))) / 8 / 1e6) Mb of LES data" for isim = 1:nsim - @info "Generating data for simulation $isim of $nsim" + # @info "Generating data for simulation $isim of $nsim" # Initial conditions u₀, p₀ = random_field(dns, T(0); pressure_solver) @@ -155,7 +156,7 @@ function create_les_data( p₀, (T(0), tburn); Δt, - processors = (step_logger(; nupdate = 10),), + # processors = (step_logger(; nupdate = 10),), pressure_solver, ) @@ -169,7 +170,7 @@ function create_les_data( Δt, processors = ( _filter_saver(_dns, _les, compression), - step_logger(; nupdate = 10), + # step_logger(; nupdate = 10), ), pressure_solver, ) @@ -178,7 +179,7 @@ function create_les_data( # Store result for current IC push!(filtered.u, f.u) # push!(filtered.p, f.p) - push!(filtered.F, f.F) + # push!(filtered.F, f.F) # push!(filtered.FG, f.FG) push!(filtered.cF, f.cF) # push!(filtered.cFG, f.cFG) diff --git a/src/closures/training.jl b/src/closures/training.jl index d7cda16a7..b0d3f0495 100644 --- a/src/closures/training.jl +++ b/src/closures/training.jl @@ -43,10 +43,8 @@ The function `loss` should take inputs like `loss(f, x, y, θ)`. The batch is moved to `device` before the loss is evaluated. """ -function create_randloss(loss, f, x, y; nuse = size(x, 2), device = identity) - x = reshape(x, size(x, 1), :) - y = reshape(y, size(y, 1), :) - nsample = size(x, 2) +function create_randloss(loss, f, x, y; nuse = 50, device = identity) + nsample = size(x)[end] d = ndims(x) function randloss(θ) i = Zygote.@ignore sort(shuffle(1:nsample)[1:nuse]) diff --git a/src/operators.jl b/src/operators.jl index f8ed914fa..edc781225 100644 --- a/src/operators.jl +++ b/src/operators.jl @@ -317,25 +317,25 @@ function laplacian!(L, p, setup) end """ - interpolate_u_p(setup, u) + interpolate_u_p(u, setup) Interpolate velocity to pressure points. """ -interpolate_u_p(setup, u) = interpolate_u_p!( - setup, +interpolate_u_p(u, setup) = interpolate_u_p!( ntuple( α -> KernelAbstractions.zeros(get_backend(u[1]), eltype(u[1]), setup.grid.N), setup.grid.dimension(), ), u, + setup, ) """ - interpolate_u_p!(setup, up, u) + interpolate_u_p!(up, u, setup) Interpolate velocity to pressure points. """ -function interpolate_u_p!(setup, up, u) +function interpolate_u_p!(up, u, setup) (; boundary_conditions, grid, Re, bodyforce) = setup (; dimension, Np, Ip) = grid D = dimension() @@ -354,12 +354,11 @@ function interpolate_u_p!(setup, up, u) end """ - interpolate_ω_p(setup, ω) + interpolate_ω_p(ω, setup) Interpolate vorticity to pressure points. """ -interpolate_ω_p(setup, ω) = interpolate_ω_p!( - setup, +interpolate_ω_p(ω, setup) = interpolate_ω_p!( setup.grid.dimension() == 2 ? KernelAbstractions.zeros(get_backend(ω), eltype(ω), setup.grid.N) : ntuple( @@ -367,16 +366,17 @@ interpolate_ω_p(setup, ω) = interpolate_ω_p!( setup.grid.dimension(), ), ω, + setup, ) """ - interpolate_ω_p!(setup, ωp, ω) + interpolate_ω_p!(ωp, ω, setup) Interpolate vorticity to pressure points. """ -interpolate_ω_p!(setup, ωp, ω) = interpolate_ω_p!(setup.grid.dimension, setup, ωp, ω) +interpolate_ω_p!(ωp, ω, setup) = interpolate_ω_p!(setup.grid.dimension, ωp, ω, setup) -function interpolate_ω_p!(::Dimension{2}, setup, ωp, ω) +function interpolate_ω_p!(::Dimension{2}, ωp, ω, setup) (; boundary_conditions, grid, Re, bodyforce) = setup (; dimension, Np, Ip) = grid D = dimension() @@ -392,7 +392,7 @@ function interpolate_ω_p!(::Dimension{2}, setup, ωp, ω) ωp end -function interpolate_ω_p!(::Dimension{3}, setup, ωp, ω) +function interpolate_ω_p!(::Dimension{3}, ωp, ω, setup) (; boundary_conditions, grid, Re) = setup (; dimension, Np, Ip) = grid D = dimension() @@ -521,10 +521,10 @@ Qfield(u, setup) = Qfield!( Compute total kinetic energy. The velocity components are interpolated to the volume centers and squared. """ -function kinetic_energy(setup, u) +function kinetic_energy(u, setup) (; dimension, Ω, Ip) = setup.grid D = dimension() - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) E = zero(eltype(up[1])) for α = 1:D # E += sum(I -> Ω[I] * up[α][I]^2, Ip) diff --git a/src/postprocess/plot_velocity.jl b/src/postprocess/plot_velocity.jl index b757e393c..260d72959 100644 --- a/src/postprocess/plot_velocity.jl +++ b/src/postprocess/plot_velocity.jl @@ -14,7 +14,7 @@ function plot_velocity(::Dimension{2}, setup, u; kwargs...) T = eltype(xp[1]) # Get velocity at pressure points - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) # qp = map((u, v) -> √(u^2 + v^2), up, vp) qp = sqrt.(up[1] .^ 2 .+ up[2] .^ 2) @@ -43,7 +43,7 @@ function plot_velocity(::Dimension{3}, setup, u; kwargs...) (; xp) = setup.grid # Get velocity at pressure points - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) qp = map((u, v, w) -> √sum(u^2 + v^2 + w^2), up...) # Levels diff --git a/src/postprocess/plot_vorticity.jl b/src/postprocess/plot_vorticity.jl index 2e8588bde..1f7cda9b8 100644 --- a/src/postprocess/plot_vorticity.jl +++ b/src/postprocess/plot_vorticity.jl @@ -18,7 +18,7 @@ function plot_vorticity(::Dimension{2}, setup, u; kwargs...) # Get fields ω = vorticity(u, setup) - ωp = interpolate_ω_p(setup, ω) + ωp = interpolate_ω_p(ω, setup) ωp = Array(ωp)[Ip] # Levels @@ -50,7 +50,7 @@ function plot_vorticity(::Dimension{3}, setup, u; kwargs...) (; grid) = setup (; xp) = grid - ωp = interpolate_ω_p(setup, vorticity(u, setup)) + ωp = interpolate_ω_p(vorticity(u, setup), setup) qp = map((u, v, w) -> √sum(u^2 + v^2 + w^2), ωp...) # Levels diff --git a/src/postprocess/save_vtk.jl b/src/postprocess/save_vtk.jl index 325f4d601..351153fdc 100644 --- a/src/postprocess/save_vtk.jl +++ b/src/postprocess/save_vtk.jl @@ -15,8 +15,8 @@ function save_vtk(setup, u, p, filename = "output/solution") D = dimension() xp = Array.(xp) vtk_grid(filename, xp...) do vtk - up = interpolate_u_p(setup, u) - ωp = interpolate_ω_p(setup, vorticity(u, setup)) + up = interpolate_u_p(u, setup) + ωp = interpolate_ω_p(vorticity(u, setup), setup) if D == 2 # ParaView prefers 3D vectors. Add zero z-component. up3 = zero(up[1]) diff --git a/src/processors/real_time_plot.jl b/src/processors/real_time_plot.jl index 14dc43ef9..1cded6595 100644 --- a/src/processors/real_time_plot.jl +++ b/src/processors/real_time_plot.jl @@ -58,10 +58,10 @@ function field_plot( (; u, p, t) = state[] _f = if fieldname == :velocity - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) elseif fieldname == :vorticity ω = vorticity(u, setup) - ωp = interpolate_ω_p(setup, ω) + ωp = interpolate_ω_p(ω, setup) elseif fieldname == :streamfunction ψ = get_streamfunction(setup, u, t) elseif fieldname == :pressure @@ -72,12 +72,12 @@ function field_plot( isnothing(sleeptime) || sleep(sleeptime) (; u, p, t) = $state f = if fieldname == :velocity - interpolate_u_p!(setup, up, u) + interpolate_u_p!(up, u, setup) map((u, v) -> √sum(u^2 + v^2), up...) elseif fieldname == :vorticity apply_bc_u!(u, t, setup) vorticity!(ω, u, setup) - interpolate_ω_p!(setup, ωp, ω) + interpolate_ω_p!(ωp, ω, setup) elseif fieldname == :streamfunction get_streamfunction!(setup, ψ, u, t) elseif fieldname == :pressure @@ -181,10 +181,10 @@ function field_plot( isnothing(sleeptime) || sleep(sleeptime) (; u, p, t) = $state f = if fieldname == :velocity - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) map((u, v, w) -> √sum(u^2 + v^2 + w^2), up...) elseif fieldname == :vorticity - ωp = interpolate_ω_p(setup, vorticity(u, setup)) + ωp = interpolate_ω_p(vorticity(u, setup), setup) map((u, v, w) -> √sum(u^2 + v^2 + w^2), ωp...) elseif fieldname == :streamfunction get_streamfunction(setup, u, t) @@ -275,7 +275,7 @@ function energy_spectrum_plot(setup, state; displayfig = true) k = Array(reshape(k, :)) ehat = @lift begin (; u, p, t) = $state - up = interpolate_u_p(setup, u) + up = interpolate_u_p(u, setup) e = sum(up -> up[Ip] .^ 2, up) Array(reshape(abs.(fft(e)[ntuple(α -> kx[α] .+ 1, D)...]), :)) end