diff --git a/docs/src/manual/differentiability.md b/docs/src/manual/differentiability.md index 9eb0027ac..8da9df70f 100644 --- a/docs/src/manual/differentiability.md +++ b/docs/src/manual/differentiability.md @@ -9,7 +9,7 @@ IncompressibleNavierStokes is which means that you can back-propagate gradients through the code. This comes at a cost however, as intermediate velocity fields need to be stored in memory for use in the backward pass. For this reason, many of the operators -come in two versions:oa slow differentiable allocating non-mutating variant (e.g. +come in two versions: a slow differentiable allocating non-mutating variant (e.g. [`divergence`](@ref)) and fast non-differentiable non-allocating mutating variant (e.g. [`divergence!`](@ref).) @@ -26,7 +26,8 @@ make a time stepping loop composed of differentiable operations: ```julia import IncompressibleNavierStokes as INS -setup = INS.Setup(; x = (0:0.01:1, 0:0.01:1), Re = 500.0) +ax = range(0, 1, 101) +setup = INS.Setup(; x = (ax, ax), Re = 500.0) psolver = INS.default_psolver(setup) method = INS.RKMethods.RK44P2() Δt = 0.01 @@ -38,7 +39,7 @@ function final_energy(u) stepper = INS.timestep(method, stepper, Δt) end (; u) = stepper - sum(abs2, u[1][Iu[1]]) / 2 + sum(abs2, u[2][Iu[2]]) / 2 + sum(abs2, u[Iu[1], 1]) / 2 + sum(abs2, u[Iu[2], 2]) / 2 end u = INS.random_field(setup) @@ -46,8 +47,7 @@ u = INS.random_field(setup) using Zygote g, = Zygote.gradient(final_energy, u) -@show size.(u) -@show size.(g) +@show size(u) size(g) ``` Now `g` is the gradient of `final_energy` with respect to the initial conditions diff --git a/docs/src/manual/time.md b/docs/src/manual/time.md index 9846c7954..28bd307cf 100644 --- a/docs/src/manual/time.md +++ b/docs/src/manual/time.md @@ -47,9 +47,6 @@ but they may be integrated in the future. ```@docs AbstractODEMethod -isexplicit -lambda_conv_max -lambda_diff_max ode_method_cache runge_kutta_method create_stepper diff --git a/examples/Kolmogorov2D.jl b/examples/Kolmogorov2D.jl index efd5cd008..5f1b5a08b 100644 --- a/examples/Kolmogorov2D.jl +++ b/examples/Kolmogorov2D.jl @@ -39,7 +39,7 @@ ustart = random_field(setup, 0.0; A = 1e-2); # # Since the force is steady, it is just stored as a field. -heatmap(setup.bodyforce[1]) +heatmap(setup.bodyforce[:, :, 1]) # ## Solve unsteady problem diff --git a/lib/Debug/perf.jl b/lib/Debug/perf.jl index f1abb1fb5..3e4a94c00 100644 --- a/lib/Debug/perf.jl +++ b/lib/Debug/perf.jl @@ -5,7 +5,7 @@ using Cthulhu using KernelAbstractions T = Float32 -n = 512 +n = 128 ax = range(T(0), T(1), n + 1) setup = Setup(; x = (ax, ax, ax), Re = T(1e3), backend = CUDABackend()) u = random_field(setup) @@ -38,18 +38,17 @@ let @benchmark IncompressibleNavierStokes.convection_adjoint!($f, $g, $u, $setup) end -f = vectorfield(setup) +f = vectorfield(setup); let fill!.(f, 0) @benchmark IncompressibleNavierStokes.convectiondiffusion!($f, $u, $setup) end -uu = stack(u) -uu = permutedims(uu, (4, 1, 2, 3)) -ff = copy(uu) +uu = stack(u); +ff = copy(uu); let fill!(ff, 0) - @benchmark IncompressibleNavierStokes.cd2!($ff, $uu, $setup) + @benchmark IncompressibleNavierStokes.arrayconvectiondiffusion!($ff, $uu, $setup) end f = vectorfield(setup) diff --git a/lib/NeuralClosure/Project.toml b/lib/NeuralClosure/Project.toml index daf079d55..491ba0510 100644 --- a/lib/NeuralClosure/Project.toml +++ b/lib/NeuralClosure/Project.toml @@ -8,6 +8,7 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" @@ -28,6 +29,7 @@ Accessors = "0.1" ComponentArrays = "0.15" DocStringExtensions = "0.9" IncompressibleNavierStokes = "2" +JLD2 = "0.5.7" KernelAbstractions = "0.9" LinearAlgebra = "1" Lux = "1" diff --git a/lib/NeuralClosure/src/NeuralClosure.jl b/lib/NeuralClosure/src/NeuralClosure.jl index 98e93aea8..1f4367c35 100644 --- a/lib/NeuralClosure/src/NeuralClosure.jl +++ b/lib/NeuralClosure/src/NeuralClosure.jl @@ -8,6 +8,7 @@ using ComponentArrays: ComponentArray using DocStringExtensions using IncompressibleNavierStokes using IncompressibleNavierStokes: Dimension, momentum!, apply_bc_u!, project! +using JLD2 using KernelAbstractions using LinearAlgebra using Lux diff --git a/lib/NeuralClosure/src/closure.jl b/lib/NeuralClosure/src/closure.jl index d7ff4cd52..07a15e8b2 100644 --- a/lib/NeuralClosure/src/closure.jl +++ b/lib/NeuralClosure/src/closure.jl @@ -2,35 +2,18 @@ Wrap closure model and parameters so that it can be used in the solver. """ function wrappedclosure(m, setup) - (; dimension, Iu) = setup.grid - D = dimension() - # function neuralclosure(u) - # u = stack(ntuple(α -> u[α][Iu[α]], D)) - # u = reshape(u, size(u)..., 1) # One sample - # mu = m(u, θ) - # mu = pad_circular(mu, 1) - # sz..., _ = size(mu) - # i = ntuple(Returns(:), D) - # mu = ntuple(α -> mu[i..., α, 1], D) - # end - neuralclosure(u, θ) = - if D == 2 - u = cat(u[1][Iu[1]], u[2][Iu[2]]; dims = 3) - u = reshape(u, size(u)..., 1) # One sample - # u = collocate(u) - mu = m(u, θ) - # mu = decollocate(mu) - mu = pad_circular(mu, 1) - mu = (mu[:, :, 1, 1], mu[:, :, 2, 1]) - elseif D == 3 - u = cat(u[1][Iu[1]], u[2][Iu[2]], u[3][Iu[3]]; dims = 4) - u = reshape(u, size(u)..., 1) # One sample - # u = collocate(u) - mu = m(u, θ) - # mu = decollocate(mu) - mu = pad_circular(mu, 1) - mu = (mu[:, :, :, 1, 1], mu[:, :, :, 2, 1], mu[:, :, :, 3, 1]) - end + (; Iu) = setup.grid + inside = Iu[1] + @assert all(==(inside), Iu) "Only periodic grids are supported" + function neuralclosure(u, θ) + s = size(u) + # u = view(u, inside, :) + u = u[inside, :] + u = reshape(u, size(u)..., 1) # Add sample dim + mu = m(u, θ) + mu = pad_circular(mu, 1) + mu = reshape(mu, s) # Remove sample dim + end end """ @@ -73,6 +56,14 @@ function collocate(u) # v = (v + circshift(v, ntuple(β -> α == β ? -1 : 0, D + 1))) / 2 # end if D == 2 + # TODO: Check if this is more efficient as + # a convolution with the two channel kernels + # [1 1; 0 0] / 2 + # and + # [0 1; 0 1] / 2 + # TODO: Maybe skip this step entirely and learn the + # collocation function subject to the skewness + # constraint (left-skewed kernel from staggered right face to center) a = selectdim(u, 3, 1) b = selectdim(u, 3, 2) a = (a .+ circshift(a, (1, 0, 0))) ./ 2 diff --git a/lib/NeuralClosure/src/cnn.jl b/lib/NeuralClosure/src/cnn.jl index 6e4047838..7d4dea2a0 100644 --- a/lib/NeuralClosure/src/cnn.jl +++ b/lib/NeuralClosure/src/cnn.jl @@ -12,7 +12,7 @@ function cnn(; rng = Random.default_rng(), ) r, c, σ, b = radii, channels, activations, use_bias - (; grid, boundary_conditions) = setup + (; grid) = setup (; dimension, x) = grid D = dimension() T = eltype(x[1]) diff --git a/lib/NeuralClosure/src/data_generation.jl b/lib/NeuralClosure/src/data_generation.jl index 726f04e56..e5ce74633 100644 --- a/lib/NeuralClosure/src/data_generation.jl +++ b/lib/NeuralClosure/src/data_generation.jl @@ -38,7 +38,7 @@ function lesdatagen(dnsobs, Φ, les, compression, psolver) FΦ = vectorfield(les) ΦF = vectorfield(les) c = vectorfield(les) - results = (; u = fill(Array.(Φu), 0), c = fill(Array.(c), 0)) + results = (; u = fill(Array(Φu), 0), c = fill(Array(c), 0)) temp = nothing on(dnsobs) do (; u, F, t) Φ(Φu, u, les, compression) @@ -47,11 +47,9 @@ function lesdatagen(dnsobs, Φ, les, compression, psolver) momentum!(FΦ, Φu, temp, t, les) apply_bc_u!(FΦ, t, les; dudt = true) project!(FΦ, les; psolver, p) - for α = 1:length(u) - c[α] .= ΦF[α] .- FΦ[α] - end - push!(results.u, Array.(Φu)) - push!(results.c, Array.(c)) + @. c = ΦF - FΦ + push!(results.u, Array(Φu)) + push!(results.c, Array(c)) end results end @@ -59,7 +57,7 @@ end """ Save filtered DNS data. """ -filtersaver( +function filtersaver( dns, les, filters, @@ -67,12 +65,12 @@ filtersaver( psolver_dns, psolver_les; nupdate = 1, + filenames, F = vectorfield(dns), p = scalarfield(dns), -) = - processor( - (results, state) -> (; results..., comptime = time() - results.comptime), - ) do state +) + @assert isnothing(filenames) || length(filenames) == length(les) * length(filters) + function initialize(state) comptime = time() t = fill(state[].t, 0) dnsobs = Observable((; state[].u, F, state[].t)) @@ -107,6 +105,20 @@ filtersaver( state[] = state[] # Save initial conditions results end + function finalize(results, state) + comptime = time() - results.comptime + (; data, t) = results + map(enumerate(data)) do (i, data) + (; u, c) = data + u = stack(u) + c = stack(c) + results = (; u, c, t, comptime) + isnothing(filenames) || jldsave(filenames[i]; results...) + results + end + end + processor(initialize, finalize) +end """ Create filtered DNS data. @@ -128,6 +140,7 @@ function create_les_data(; icfunc = (setup, psolver, rng) -> random_field(setup, typeof(Re)(0); psolver, rng), processors = (; log = timelogger(; nupdate = 10)), rng, + filenames = nothing, kwargs..., ) T = typeof(Re) @@ -196,12 +209,13 @@ function create_les_data(; psolver, psolver_les; nupdate = savefreq, + filenames, # Reuse arrays from cache to save memory in 3D DNS. # Since processors are called outside # Runge-Kutta steps, there is no danger # in overwriting the arrays. - F = cache.u₀, + F = cache.ustart, p = cache.p, ), ), @@ -226,7 +240,7 @@ function create_io_arrays(data, setup) for it = 1:nt, α = 1:D copyto!( view(u, colons..., α, it), - view(getfield(trajectory, usym)[it][α], Iu[α]), + view(getfield(trajectory, usym), Iu[α], α, it), ) end u diff --git a/lib/NeuralClosure/src/filter.jl b/lib/NeuralClosure/src/filter.jl index 8d6691e07..add5586a2 100644 --- a/lib/NeuralClosure/src/filter.jl +++ b/lib/NeuralClosure/src/filter.jl @@ -24,34 +24,33 @@ struct VolumeAverage <: AbstractFilter end Φ(vectorfield(setup_les), u, setup_les, compression) function (::FaceAverage)(v, u, setup_les, comp) - (; grid, workgroupsize) = setup_les - (; Nu, Iu) = grid - D = length(u) + (; grid, backend, workgroupsize) = setup_les + (; dimension, Nu, Iu) = grid + D = dimension() @kernel function Φ!(v, u, ::Val{α}, face, I0) where {α} I = @index(Global, Cartesian) J = I0 + comp * (I - oneunit(I)) - s = zero(eltype(v[α])) + s = zero(eltype(v)) for i in face - s += u[α][J+i] + s += u[J+i, α] end - v[α][I0+I] = s / comp^(D - 1) + v[I0+I, α] = s / comp^(D - 1) end for α = 1:D ndrange = Nu[α] - I0 = first(Iu[α]) - I0 -= oneunit(I0) + I0 = getoffset(Iu[α]) face = CartesianIndices(ntuple(β -> β == α ? (comp:comp) : (1:comp), D)) - Φ!(get_backend(v[1]), workgroupsize)(v, u, Val(α), face, I0; ndrange) + Φ!(backend, workgroupsize)(v, u, Val(α), face, I0; ndrange) end v end "Reconstruct DNS velocity `u` from LES velocity `v`." function reconstruct!(u, v, setup_dns, setup_les, comp) - (; grid, boundary_conditions, workgroupsize) = setup_les - (; N, Iu) = grid - D = length(u) - e = Offset{D}() + (; grid, boundary_conditions, backend, workgroupsize) = setup_les + (; dimension, N) = grid + D = dimension() + e = Offset(D) @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions) @kernel function R!(u, v, ::Val{α}, volume) where {α} J = @index(Global, Cartesian) @@ -61,15 +60,15 @@ function reconstruct!(u, v, setup_dns, setup_les, comp) Jleft.I[α] == 1 && (Jleft += (N[α] - 2) * e(α)) for i in volume s = zero(eltype(v[α])) - s += (comp - i.I[α]) * v[α][J] - s += i.I[α] * v[α][Jleft] - u[α][I-i] = s / comp + s += (comp - i.I[α]) * v[J, α] + s += i.I[α] * v[Jleft, α] + u[I-i, α] = s / comp end end for α = 1:D ndrange = N .- 2 volume = CartesianIndices(ntuple(β -> 0:comp-1, D)) - R!(get_backend(v[1]), workgroupsize)(u, v, Val(α), volume; ndrange) + R!(backend, workgroupsize)(u, v, Val(α), volume; ndrange) end u end @@ -79,30 +78,29 @@ reconstruct(v, setup_dns, setup_les, comp) = reconstruct!(vectorfield(setup_dns), v, setup_dns, setup_les, comp) function (::VolumeAverage)(v, u, setup_les, comp) - (; grid, boundary_conditions, workgroupsize) = setup_les - (; N, Nu, Iu) = grid - D = length(u) + (; grid, boundary_conditions, backend, workgroupsize) = setup_les + (; dimension, N, Nu, Iu) = grid + D = dimension() @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions) @kernel function Φ!(v, u, ::Val{α}, volume, I0) where {α} I = @index(Global, Cartesian) J = I0 + comp * (I - oneunit(I)) - s = zero(eltype(v[α])) + s = zero(eltype(v)) # n = 0 for i in volume # Periodic extension K = J + i K = mod1.(K.I, comp .* (N .- 2)) K = CartesianIndex(K) - s += u[α][K] + s += u[K, α] # n += 1 end n = (iseven(comp) ? comp + 1 : comp) * comp^(D - 1) - v[α][I0+I] = s / n + v[I0+I, α] = s / n end for α = 1:D ndrange = Nu[α] - I0 = first(Iu[α]) - I0 -= oneunit(I0) + I0 = getoffset(Iu[α]) volume = CartesianIndices( ntuple( β -> @@ -112,7 +110,7 @@ function (::VolumeAverage)(v, u, setup_les, comp) D, ), ) - Φ!(get_backend(v[1]), workgroupsize)(v, u, Val(α), volume, I0; ndrange) + Φ!(backend, workgroupsize)(v, u, Val(α), volume, I0; ndrange) end v end diff --git a/lib/NeuralClosure/src/groupconv.jl b/lib/NeuralClosure/src/groupconv.jl index 4d04f09df..8e201a078 100644 --- a/lib/NeuralClosure/src/groupconv.jl +++ b/lib/NeuralClosure/src/groupconv.jl @@ -41,20 +41,42 @@ function rot2(u, r) u[I, chans...] end -# For vector fields (u, v) +"Rotate vector fields `[ux;;; uy]`" function rot2(u::Tuple{T,T}, r) where {T} + # ux, uy = eachslice(u; dims = ndims(u)) + ux, uy = u r = mod(r, 4) - ru = rot2(u[1], r) - rv = rot2(u[2], r) - if r == 0 - (ru, rv) + rx = rot2(ux, r) + ry = rot2(uy, r) + ru = if r == 0 + (rx, ry) + elseif r == 1 + (-ry, rx) + elseif r == 2 + (-rx, -ry) + elseif r == 3 + (ry, -rx) + end + ru +end + +"Rotate vector fields `[ux;;; uy]`" +function vecrot2(u, r) + # ux, uy = eachslice(u; dims = ndims(u)) + ux, uy = u[:, :, 1], u[:, :, 2] + r = mod(r, 4) + rx = rot2(ux, r) + ry = rot2(uy, r) + ru = if r == 0 + (rx, ry) elseif r == 1 - (-rv, ru) + (-ry, rx) elseif r == 2 - (-ru, -rv) + (-rx, -ry) elseif r == 3 - (rv, -ru) + (ry, -rx) end + stack(ru) end # # For augmented vector fields (u, v, -u, -v) @@ -77,8 +99,9 @@ end "Rotate staggered grid velocity field. See also [`rot2`](@ref)." function rot2stag(u, g) g = mod(g, 4) - u = rot2(u, g) - ux, uy = u + u = vecrot2(u, g) + # ux, uy = eachslice(u; dims = ndims(u)) + ux, uy = u[:, :, 1], u[:, :, 2] if g in (1, 2) ux = circshift(ux, -1) ux[end, :] .= ux[2, :] @@ -87,7 +110,7 @@ function rot2stag(u, g) uy = circshift(uy, (0, -1)) uy[:, end] .= uy[:, 2] end - (ux, uy) + cat(ux, uy; dims = 3) end """ diff --git a/lib/NeuralClosure/src/training.jl b/lib/NeuralClosure/src/training.jl index 014390303..a33019b91 100644 --- a/lib/NeuralClosure/src/training.jl +++ b/lib/NeuralClosure/src/training.jl @@ -32,7 +32,8 @@ create_dataloader_post(trajectories; ntrajectory, nunroll, device = identity) = @assert nt ≥ nunroll "Trajectory too short for nunroll = $nunroll" istart = rand(rng, 1:nt-nunroll) it = istart:istart+nunroll - (; u = device.(u[it]), t = t[it]) + u = selectdim(u, ndims(u), it) |> Array |> device # convert view to array first + (; u, t = t[it]) end data, rng end @@ -115,25 +116,25 @@ Create a-posteriori loss function. function create_loss_post(; setup, method, psolver, closure, nsubstep = 1) closure_model = wrappedclosure(closure, setup) setup = (; setup..., closure_model) - (; dimension, Iu, x) = setup.grid - D = dimension() + (; Iu) = setup.grid + inside = Iu[1] + @assert all(==(inside), Iu) loss_post(data, θ) = sum(data) do (; u, t) T = eltype(θ) - v = u[1] + ules = selectdim(u, ndims(u), 1) |> collect stepper = - create_stepper(method; setup, psolver, u = v, temp = nothing, t = t[1]) + create_stepper(method; setup, psolver, u = ules, temp = nothing, t = t[1]) loss = zero(T) for it = 2:length(t) Δt = (t[it] - t[it-1]) / nsubstep for isub = 1:nsubstep stepper = timestep(method, stepper, Δt; θ) end - a, b = T(0), T(0) - for α = 1:length(u[1]) - a += sum(abs2, (stepper.u[α]-u[it][α])[Iu[α]]) - b += sum(abs2, u[it][α][Iu[α]]) - end + uref = view(u, inside, :, it) + ules = view(stepper.u, inside, :) + a = sum(abs2, ules - uref) + b = sum(abs2, uref) loss += a / b end loss / (length(t) - 1) @@ -145,14 +146,15 @@ Create a-posteriori relative error. """ function create_relerr_post(; data, setup, method, psolver, closure_model, nsubstep = 1) setup = (; setup..., closure_model) - (; dimension, Iu) = setup.grid - D = dimension() + (; Iu) = setup.grid + inside = Iu[1] + @assert all(==(inside), Iu) (; u, t) = data - v = copy.(u[1]) + v = selectdim(u, ndims(u), 1) |> collect cache = IncompressibleNavierStokes.ode_method_cache(method, setup) function relerr_post(θ) - T = eltype(u[1][1]) - copyto!.(v, u[1]) + T = eltype(u) + copyto!(v, selectdim(u, ndims(u), 1)) stepper = create_stepper(method; setup, psolver, u = v, temp = nothing, t = t[1]) e = zero(T) for it = 2:length(t) @@ -161,13 +163,10 @@ function create_relerr_post(; data, setup, method, psolver, closure_model, nsubs stepper = IncompressibleNavierStokes.timestep!(method, stepper, Δt; θ, cache) end - a, b = T(0), T(0) - for α = 1:D - # a += sum(abs2, (stepper.u[α]-u[it][α])[Iu[α]]) - # b += sum(abs2, u[it][α][Iu[α]]) - a += sum(abs2, view(stepper.u[α] - u[it][α], Iu[α])) - b += sum(abs2, view(u[it][α], Iu[α])) - end + uref = view(u, inside, :, it) + ules = view(stepper.u, inside, :) + a = sum(abs2, ules - uref) + b = sum(abs2, uref) e += sqrt(a) / sqrt(b) end e / (length(t) - 1) @@ -189,15 +188,17 @@ function create_relerr_symmetry_post(; (; dimension, Iu) = setup.grid D = dimension() T = eltype(u[1]) + inside = Iu[1] + @assert all(==(inside), Iu) cache = IncompressibleNavierStokes.ode_method_cache(method, setup) function err(θ) stepper = - create_stepper(method; setup, psolver, u = copy.(u), temp = nothing, t = T(0)) + create_stepper(method; setup, psolver, u = copy(u), temp = nothing, t = T(0)) stepper_rot = create_stepper( method; setup, psolver, - u = rot2stag(copy.(u), g), + u = rot2stag(copy(u), g), temp = nothing, t = T(0), ) @@ -207,11 +208,8 @@ function create_relerr_symmetry_post(; stepper_rot = IncompressibleNavierStokes.timestep!(method, stepper_rot, Δt; θ, cache) u_rot = rot2stag(stepper.u, g) - a, b = T(0), T(0) - for α = 1:D - a += sum(abs2, view(stepper_rot.u[α] - u_rot[α], Iu[α])) - b += sum(abs2, view(u_rot[α], Iu[α])) - end + a = sum(abs2, view(stepper_rot.u - u_rot, inside, :)) + b = sum(abs2, view(u_rot, inside, :)) e += sqrt(a) / sqrt(b) end e / nstep @@ -225,16 +223,17 @@ function create_relerr_symmetry_prior(; u, setup, g = 1) (; grid, closure_model) = setup (; dimension, Iu) = grid D = dimension() - T = eltype(u[1][1]) + T = eltype(u[1]) + inside = Iu[1] + @assert all(==(inside), Iu) function err(θ) - e = sum(u) do u + e = sum(eachslice(u; dims = ndims(u))) do u cr = closure_model(rot2stag(u, g), θ) rc = rot2stag(closure_model(u, θ), g) - a, b = T(0), T(0) - for α = 1:D - a += sum(abs2, view(rc[α] - cr[α], Iu[α])) - b += sum(abs2, view(cr[α], Iu[α])) - end + cr = view(cr, inside, :) + rc = view(rc, inside, :) + a = sum(abs2, rc - cr) + b = sum(abs2, cr) sqrt(a) / sqrt(b) end e / length(u) @@ -306,6 +305,6 @@ function create_callback( (; callbackstate, callback) end +getlearningrate(r) = -1 # Fallback getlearningrate(r::Adam) = r.eta getlearningrate(r::OptimiserChain{Tuple{Adam,WeightDecay}}) = r.opts[1].eta -getlearningrate(r) = -1 diff --git a/lib/NeuralClosure/test/examplerun.jl b/lib/NeuralClosure/test/examplerun.jl index a5f15f642..82bc38f64 100644 --- a/lib/NeuralClosure/test/examplerun.jl +++ b/lib/NeuralClosure/test/examplerun.jl @@ -19,12 +19,8 @@ filters = (FaceAverage(), VolumeAverage()), ) - data = stack(splitseed(123, 3)) do seed - (; data, t, comptime) = create_les_data(; params..., rng = Xoshiro(seed)) - map(data) do (; u, c) - (; u, c, t, comptime) - end - end + data = + stack(seed -> create_les_data(; params..., rng = Xoshiro(seed)), splitseed(123, 3)) # Build LES setups and assemble operators setups = map( @@ -37,9 +33,10 @@ # Create input/output arrays for a-priori training (ubar vs c) io = map( - I -> create_io_arrays(data[I[1], I[2], :], setups[I[1]]), Iterators.product(eachindex(params.nles), eachindex(params.filters)), - ) + ) do (igrid, ifil) + create_io_arrays(data[igrid, ifil, :], setups[igrid]) + end m_cnn = let rng = Xoshiro(123) @@ -128,7 +125,7 @@ optstate = Optimisers.setup(opt, θ) d = data[ig, ifil, 1] it = 1:5 - snap = (; u = d.u[it], t = d.t[it]) + snap = (; u = selectdim(d.u, ndims(d.u), it), t = d.t[it]) (; callbackstate, callback) = create_callback( create_relerr_post(; data = snap, @@ -175,7 +172,7 @@ psolver = psolver_spectral(setup) traj = data[ig, ifil, itraj] it = 1:5 - snaps = (; u = traj.u[it], t = traj.t[it]) + snaps = (; u = selectdim(traj.u, ndims(traj.u), it), t = traj.t[it]) for (im, m) in enumerate(models) err = create_relerr_post(; data = snaps, @@ -212,7 +209,7 @@ setup = (; setup..., closure_model = wrappedclosure(m.closure, setup)) traj = data[ig, ifil, itraj] err = create_relerr_symmetry_post(; - u = traj.u[1], + u = selectdim(traj.u, ndims(traj.u), 1), setup, psolver = psolver_spectral(setup), Δt = (traj.t[2] - traj.t[1]), diff --git a/lib/NeuralClosure/test/filter.jl b/lib/NeuralClosure/test/filter.jl index 4d9ddaa4a..32629f6fa 100644 --- a/lib/NeuralClosure/test/filter.jl +++ b/lib/NeuralClosure/test/filter.jl @@ -5,14 +5,14 @@ setups = map(n -> Setup(; x = ntuple(d -> range(0.0, 1.0, n + 1), D), Re = 1e3), n) u = vectorfield.(setups) val = 3.83 - fill!.(u[end], val) + fill!(u[end], val) compression = div.(n[end], n) filters = FaceAverage(), VolumeAverage() for Φ in filters, (setup, comp) in zip(setups, compression) (; Iu) = setup.grid v = Φ(u[end], setup, comp) for i = 1:D - @test all(≈(val), v[i][Iu[i]]) + @test all(≈(val), v[Iu[i], i]) end end end diff --git a/lib/PaperDC/postanalysis.jl b/lib/PaperDC/postanalysis.jl index 3bc025fb9..c4ca92fec 100644 --- a/lib/PaperDC/postanalysis.jl +++ b/lib/PaperDC/postanalysis.jl @@ -512,7 +512,10 @@ let psolver = psolver_spectral(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) it = 1:100 - data = (; u = device.(sample.u[it]), t = sample.t[it]) + data = (; + u = selectdim(sample.u, ndims(sample.u), it) |> collect |> device, + t = sample.t[it], + ) nsubstep = 5 method = RKProject(params.method, projectorder) ## No model @@ -677,8 +680,8 @@ let setup = getsetup(; params, nles) psolver = default_psolver(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) - ustart = sample.u[1] |> device - T = eltype(ustart[1]) + ustart = selectdim(sample.u, ndims(sample.u), 1) |> collect |> device + T = eltype(ustart) # Shorter time for DIF t_DIF = T(1) @@ -687,8 +690,11 @@ let divergencehistory.ref[I] = let div = scalarfield(setup) udev = vectorfield(setup) - map(sample.t[1:5:end], sample.u[1:5:end]) do t, u - copyto!.(udev, u) + it = 1:5:length(sample.t) + map(it) do it + t = sample.t[it] + u = selectdim(sample.u, ndims(sample.u), it) + copyto!(udev, u) IncompressibleNavierStokes.divergence!(div, udev, setup) d = view(div, setup.grid.Ip) d = sum(abs2, d) / length(d) @@ -696,11 +702,16 @@ let Point2f(t, d) end end - energyhistory.ref[I] = map( - (t, u) -> Point2f(t, total_kinetic_energy(device(u), setup)), - sample.t[1:5:end], - sample.u[1:5:end], - ) + energyhistory.ref[I] = let + it = 1:5:length(sample.t) + udev = vectorfield(setup) + map(it) do it + t = sample.t[it] + u = selectdim(sample.u, ndims(sample.u), it) + copyto!(udev, u) + Point2f(t, total_kinetic_energy(udev, setup)) + end + end nupdate = 5 writer = processor() do state @@ -948,7 +959,7 @@ end let s = length(params.nles), length(params.filters), length(projectorders) - temp = ntuple(Returns(zeros(T, ntuple(Returns(0), params.D))), params.D) + temp = zeros(T, ntuple(Returns(0), params.D + 1)) keys = [:ref, :nomodel, :smag, :cnn_prior, :cnn_post] times = T[0.1, 0.5, 1.0, 5.0] itime_max_DIF = 3 @@ -963,7 +974,7 @@ let setup = getsetup(; params, nles) psolver = default_psolver(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) - ustart = sample.u[1] + ustart = selectdim(sample.u, ndims(sample.u), 1) |> collect t = sample.t solve(ustart, tlims, closure_model, θ) = solve_unsteady(; @@ -996,7 +1007,7 @@ let getprev(i, sym) = i == 1 ? ustart : utimes[i-1][sym][I] # Compute fields - utimes[i].ref[I] = sample.u[it] + utimes[i].ref[I] = selectdim(sample.u, ndims(sample.u), it) |> collect utimes[i].nomodel[I] = solve(getprev(i, :nomodel), tlims, nothing, nothing) utimes[i].smag[I] = solve(getprev(i, :smag), tlims, smagorinsky_closure(setup), θ_smag[I]) diff --git a/lib/PaperDC/postanalysis3D.jl b/lib/PaperDC/postanalysis3D.jl index 32b92a365..eda872585 100644 --- a/lib/PaperDC/postanalysis3D.jl +++ b/lib/PaperDC/postanalysis3D.jl @@ -539,7 +539,10 @@ let psolver = psolver_spectral(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) it = 1:20 - data = (; u = device.(sample.u[it]), t = sample.t[it]) + data = (; + u = selectdim(sample.u, ndims(sample.u), it) |> collect |> device, + t = sample.t[it], + ) nsubstep = 5 method = RKProject(params.method, projectorder) ## No model @@ -703,7 +706,7 @@ let setup = getsetup(; params, nles) psolver = default_psolver(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) - ustart = sample.u[1] |> device + ustart = selectdim(sample.u, ndims(sample.u), 1) |> collect |> device T = eltype(ustart[1]) # Shorter time for DIF @@ -714,8 +717,11 @@ let divergencehistory.ref[I] = let div = scalarfield(setup) udev = vectorfield(setup) - map(sample.t[1:nupdate:end], sample.u[1:nupdate:end]) do t, u - copyto!.(udev, u) + it = 1:nupdate:length(sample.t) + map(it) do it + t = sample.t[it] + u = selectdim(sample.u, ndims(sample.u), it) + copyto!(udev, u) IncompressibleNavierStokes.divergence!(div, udev, setup) d = view(div, setup.grid.Ip) d = sum(abs2, d) / length(d) @@ -723,11 +729,16 @@ let Point2f(t, d) end end - energyhistory.ref[I] = map( - (t, u) -> Point2f(t, total_kinetic_energy(device(u), setup)), - sample.t[1:nupdate:end], - sample.u[1:nupdate:end], - ) + energyhistory.ref[I] = let + it = 1:nupdate:length(sample.t) + udev = vectorfield(setup) + map(it) do it + t = sample.t[it] + u = selectdim(sample.u, ndims(sample.u), it) + copyto!(udev, u) + Point2f(t, total_kinetic_energy(udev, setup)) + end + end writer = processor() do state div = scalarfield(setup) @@ -965,7 +976,7 @@ end let s = length(params.nles), length(params.filters), length(projectorders) - temp = ntuple(Returns(zeros(T, ntuple(Returns(0), params.D))), params.D) + temp = zeros(T, ntuple(Returns(0), params.D + 1)) keys = [:ref, :nomodel, :smag, :cnn_prior, :cnn_post] times = T[0.1, 0.5, 1.0, 3.0] itime_max_DIF = 3 @@ -981,6 +992,7 @@ let psolver = default_psolver(setup) sample = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_test[1])) ustart = sample.u[1] + ustart = selectdim(sample.u, ndims(sample.u), 1) |> collect t = sample.t solve(ustart, tlims, closure_model, θ) = solve_unsteady(; @@ -1013,7 +1025,7 @@ let getprev(i, sym) = i == 1 ? ustart : utimes[i-1][sym][I] # Compute fields - utimes[i].ref[I] = sample.u[it] + utimes[i].ref[I] = selectdim(sample.u, ndims(sample.u), it) |> collect utimes[i].nomodel[I] = solve(getprev(i, :nomodel), tlims, nothing, nothing) utimes[i].smag[I] = solve(getprev(i, :smag), tlims, smagorinsky_closure(setup), θ_smag[I]) diff --git a/lib/PaperDC/prioranalysis.jl b/lib/PaperDC/prioranalysis.jl index 213622cee..5235c6331 100644 --- a/lib/PaperDC/prioranalysis.jl +++ b/lib/PaperDC/prioranalysis.jl @@ -201,7 +201,7 @@ case.D == 2 && with_theme() do apply_bc_u!(PF, T(0), dns.setup) ΦPF = fil.Φ(PF, fil.setup, fil.compression) apply_bc_u!(ΦPF, T(0), fil.setup) - c = ΦPF .- PFv + c = ΦPF - PFv apply_bc_u!(c, T(0), fil.setup) ## Make plots diff --git a/lib/PaperDC/src/observe.jl b/lib/PaperDC/src/observe.jl index 7d84a83d9..05ba965a6 100644 --- a/lib/PaperDC/src/observe.jl +++ b/lib/PaperDC/src/observe.jl @@ -3,14 +3,14 @@ function observe_v(dnsobs, Φ, les, compression, psolver) (; dimension, N, Iu, Ip) = grid D = dimension() Mα = N[1] - 2 - v = zero.(Φ(dnsobs[].u, les, compression)) - Pv = zero.(v) - p = zero(v[1]) - div = zero(p) - ΦPF = zero.(v) - PFΦ = zero.(v) - c = zero.(v) - T = eltype(v[1]) + v = vectorfield(les) + Pv = vectorfield(les) + p = scalarfield(les) + div = scalarfield(les) + ΦPF = vectorfield(les) + PFΦ = vectorfield(les) + c = vectorfield(les) + T = eltype(v) results = (; Φ, Mα, @@ -30,28 +30,28 @@ function observe_v(dnsobs, Φ, les, compression, psolver) momentum!(PFΦ, v, nothing, t, les) apply_bc_u!(PFΦ, t, les; dudt = true) project!(PFΦ, les; psolver, p) - foreach(α -> c[α] .= ΦPF[α] .- PFΦ[α], 1:D) + @. c = ΦPF - PFΦ apply_bc_u!(c, t, les) divergence!(div, v, les) norm_Du = norm(div[Ip]) - norm_v = sqrt(sum(α -> sum(abs2, v[α][Iu[α]]), 1:D)) + norm_v = sqrt(sum(α -> sum(abs2, v[Iu[α], α]), 1:D)) push!(results.Dv, norm_Du / norm_v) - copyto!.(Pv, v) + copyto!(Pv, v) project!(Pv, les; psolver, p) - foreach(α -> Pv[α] .= Pv[α] .- v[α], 1:D) - norm_vmPv = sqrt(sum(α -> sum(abs2, Pv[α][Iu[α]]), 1:D)) + @. Pv = Pv - v + norm_vmPv = sqrt(sum(α -> sum(abs2, Pv[Iu[α], α]), 1:D)) push!(results.Pv, norm_vmPv / norm_v) Pc = Pv - copyto!.(Pc, c) + copyto!(Pc, c) project!(Pc, les; psolver, p) - foreach(α -> Pc[α] .= Pc[α] .- c[α], 1:D) - norm_cmPc = sqrt(sum(α -> sum(abs2, Pc[α][Iu[α]]), 1:D)) - norm_c = sqrt(sum(α -> sum(abs2, c[α][Iu[α]]), 1:D)) + @. Pc = Pc - c + norm_cmPc = sqrt(sum(α -> sum(abs2, Pc[Iu[α], α]), 1:D)) + norm_c = sqrt(sum(α -> sum(abs2, c[Iu[α], α]), 1:D)) push!(results.Pc, norm_cmPc / norm_c) - norm_ΦPF = sqrt(sum(α -> sum(abs2, ΦPF[α][Iu[α]]), 1:D)) + norm_ΦPF = sqrt(sum(α -> sum(abs2, ΦPF[Iu[α], α]), 1:D)) push!(results.c, norm_c / norm_ΦPF) kinetic_energy!(p, v, les) diff --git a/lib/PaperDC/src/rk.jl b/lib/PaperDC/src/rk.jl index 6c16cbbc5..675308e71 100644 --- a/lib/PaperDC/src/rk.jl +++ b/lib/PaperDC/src/rk.jl @@ -46,19 +46,17 @@ function IncompressibleNavierStokes.timestep!( cache, ) (; setup, psolver, u, temp, t, n) = stepper - (; grid, closure_model) = setup - (; dimension, Iu) = grid + (; closure_model) = setup (; rk, projectorder) = method (; A, b, c) = rk - (; u₀, ku, p) = cache - D = dimension() + (; ustart, ku, p) = cache nstage = length(b) m = closure_model @assert projectorder ∈ instances(ProjectOrder.T) "Unknown projectorder: $projectorder" # Update current solution - t₀ = t - copyto!.(u₀, u) + tstart = t + copyto!(ustart, u) for i = 1:nstage # Compute force at current stage i @@ -72,7 +70,7 @@ function IncompressibleNavierStokes.timestep!( end # Add closure term - isnothing(m) || map((k, m) -> k .+= m, ku[i], m(u, θ)) + isnothing(m) || (ku[i] .+= m(u, θ)) # Project F second if projectorder == ProjectOrder.Second @@ -81,15 +79,12 @@ function IncompressibleNavierStokes.timestep!( end # Intermediate time step - t = t₀ + c[i] * Δt + t = tstart + c[i] * Δt # Apply stage forces - for α = 1:D - u[α] .= u₀[α] - for j = 1:i - @. u[α] += Δt * A[i, j] * ku[j][α] - # @. u[α][Iu[α]] += Δt * A[i, j] * ku[j][α][Iu[α]] - end + u .= ustart + for j = 1:i + @. u += Δt * A[i, j] * ku[j] end # Project stage u directly @@ -110,18 +105,16 @@ end function IncompressibleNavierStokes.timestep(method::RKProject, stepper, Δt; θ = nothing) (; setup, psolver, u, temp, t, n) = stepper - (; grid, closure_model) = setup - (; dimension) = grid + (; closure_model) = setup (; rk, projectorder) = method (; A, b, c) = rk - D = dimension() nstage = length(b) m = closure_model @assert projectorder ∈ instances(ProjectOrder.T) "Unknown projectorder: $projectorder" # Update current solution (does not depend on previous step size) - t₀ = t - u₀ = u + tstart = t + ustart = u ku = () for i = 1:nstage @@ -136,7 +129,7 @@ function IncompressibleNavierStokes.timestep(method::RKProject, stepper, Δt; θ end # Add closure term - isnothing(m) || (F = F .+ m(u, θ)) + isnothing(m) || (F = F + m(u, θ)) # Project F second if projectorder == ProjectOrder.Second @@ -148,13 +141,12 @@ function IncompressibleNavierStokes.timestep(method::RKProject, stepper, Δt; θ ku = (ku..., F) # Intermediate time step - t = t₀ + c[i] * Δt + t = tstart + c[i] * Δt # Apply stage forces - u = u₀ + u = ustart for j = 1:i u = @. u + Δt * A[i, j] * ku[j] - # u = tupleadd(u, @.(Δt * A[i, j] * ku[j])) end # Project stage u directly diff --git a/lib/PaperDC/src/train.jl b/lib/PaperDC/src/train.jl index 9233f6f3f..9f3752ff1 100644 --- a/lib/PaperDC/src/train.jl +++ b/lib/PaperDC/src/train.jl @@ -11,18 +11,19 @@ createdata(; params, seeds, outdir, taskid) = @info "Skipping seed $(repr(seed)) for task $taskid" continue end - (; data, t, comptime) = create_les_data(; params..., rng = Xoshiro(seed)) - @info("Trajectory info:", comptime / 60, length(t), Base.summarysize(data) * 1e-9,) - for (ifilter, Φ) in enumerate(params.filters), - (igrid, nles) in enumerate(params.nles) - - (; u, c) = data[igrid, ifilter] + filenames = map(Iterators.product(params.nles, params.filters)) do nles, Φ f = getdatafile(outdir, nles, Φ, seed) datadir = dirname(f) ispath(datadir) || mkpath(datadir) - @info "Saving data to $f" - jldsave(f; u, c, t, comptime) + f end + data = create_les_data(; params..., rng = Xoshiro(seed), filenames) + @info( + "Trajectory info:", + data[1].comptime / 60, + length(data[1].t), + Base.summarysize(data) * 1e-9, + ) end getpriorfile(outdir, nles, filter) = @@ -236,7 +237,8 @@ function trainpost(; (; callbackstate, callback) = let d = data_valid[1] it = 1:nunroll_valid - data = (; u = device.(d.u[it]), t = d.t[it]) + data = + (; u = selectdim(d.u, ndims(d.u), it) |> collect |> device, t = d.t[it]) create_callback( create_relerr_post(; data, @@ -322,7 +324,7 @@ function trainsmagorinsky(; psolver = default_psolver(setup) d = namedtupleload(getdatafile(outdir, nles, Φ, dns_seeds_train[1])) it = 1:nunroll - data = (; u = device.(d.u[it]), t = d.t[it]) + data = (; u = selectdim(d.u, ndims(d.u), it) |> collect |> device, t = d.t[it]) θmin = T(0) emin = T(Inf) err = create_relerr_post(; diff --git a/src/IncompressibleNavierStokes.jl b/src/IncompressibleNavierStokes.jl index 5469a1246..684ac04ca 100644 --- a/src/IncompressibleNavierStokes.jl +++ b/src/IncompressibleNavierStokes.jl @@ -73,8 +73,6 @@ include("utils.jl") include("time_steppers/methods.jl") include("time_steppers/time_stepper_caches.jl") include("time_steppers/step.jl") -include("time_steppers/isexplicit.jl") -include("time_steppers/lambda_max.jl") include("time_steppers/RKMethods.jl") # Precompile workflow @@ -111,7 +109,7 @@ export solve_unsteady, timestep, create_stepper export velocityfield, temperaturefield, random_field # Utils -export splitseed, plotgrid, save_vtk, get_lims +export getoffset, splitseed, plotgrid, save_vtk, get_lims # ODE methods export AdamsBashforthCrankNicolsonMethod, OneLegMethod, RKMethods, LMWray3 diff --git a/src/boundary_conditions.jl b/src/boundary_conditions.jl index 7aa30f04e..57a30fefa 100644 --- a/src/boundary_conditions.jl +++ b/src/boundary_conditions.jl @@ -103,7 +103,7 @@ function boundary(β, N, I, isright) end "Apply velocity boundary conditions (differentiable version)." -apply_bc_u(u, t, setup; kwargs...) = apply_bc_u!(copy.(u), t, setup; kwargs...) +apply_bc_u(u, t, setup; kwargs...) = apply_bc_u!(copy(u), t, setup; kwargs...) "Apply pressure boundary conditions (differentiable version)." apply_bc_p(p, t, setup; kwargs...) = apply_bc_p!(copy(p), t, setup; kwargs...) @@ -118,7 +118,7 @@ ChainRulesCore.rrule(::typeof(apply_bc_u), u, t, setup; kwargs...) = ( NoTangent(), # Important: identity operator should be part of `apply_bc_u_pullback`, # but is actually implemented via the `copy` below instead. - Tangent{typeof(u)}(apply_bc_u_pullback!(copy.((φbar...,)), t, setup; kwargs...)...), + apply_bc_u_pullback!(copy(unthunk(φbar)), t, setup; kwargs...), NoTangent(), NoTangent(), ), @@ -159,7 +159,7 @@ ChainRulesCore.rrule(::typeof(apply_bc_temp), temp, t, setup) = ( "Apply velocity boundary conditions (in-place version)." function apply_bc_u!(u, t, setup; kwargs...) (; boundary_conditions) = setup - D = length(u) + D = setup.grid.dimension() for β = 1:D apply_bc_u!(boundary_conditions[β][1], u, β, t, setup; isright = false, kwargs...) apply_bc_u!(boundary_conditions[β][2], u, β, t, setup; isright = true, kwargs...) @@ -278,15 +278,13 @@ function apply_bc_u!(::PeriodicBC, u, β, t, setup; isright, kwargs...) isright && return u # We do both in one go for "left" (; dimension, N, Iu) = setup.grid D = dimension() - for α = 1:D - uα, eβ = u[α], Offset{D}()(β) - Ia = boundary(β, N, Iu[α], false) - Ib = boundary(β, N, Iu[α], true) - Ja = Ia .+ eβ - Jb = Ib .- eβ - @. uα[Ia] = uα[Jb] - @. uα[Ib] = uα[Ja] - end + eβ = Offset(D)(β) + Ia = boundary(β, N, Iu[1], false) + Ib = boundary(β, N, Iu[1], true) + Ja = Ia .+ eβ + Jb = Ib .- eβ + @. u[Ia, :] = u[Jb, :] + @. u[Ib, :] = u[Ja, :] u end @@ -294,17 +292,15 @@ function apply_bc_u_pullback!(::PeriodicBC, φbar, β, t, setup; isright, kwargs isright && return φbar # We do both in one go for "left" (; dimension, N, Iu) = setup.grid D = dimension() - for α = 1:D - φα, eβ = φbar[α], Offset{D}()(β) - Ia = boundary(β, N, Iu[α], false) - Ib = boundary(β, N, Iu[α], true) - Ja = Ia .+ eβ - Jb = Ib .- eβ - @. φα[Jb] += φα[Ia] - @. φα[Ja] += φα[Ib] - @. φα[Ia] = 0 - @. φα[Ib] = 0 - end + eβ = Offset(D)(β) + Ia = boundary(β, N, Iu[1], false) + Ib = boundary(β, N, Iu[1], true) + Ja = Ia .+ eβ + Jb = Ib .- eβ + @. φbar[Jb, :] += φbar[Ia, :] + @. φbar[Ja, :] += φbar[Ib, :] + @. φbar[Ia, :] = 0 + @. φbar[Ib, :] = 0 φbar end @@ -312,7 +308,7 @@ function apply_bc_p!(bc::PeriodicBC, p, β, t, setup; isright, kwargs...) isright && return p # We do both in one go for "left" (; dimension, N, Ip) = setup.grid D = dimension() - eβ = Offset{D}()(β) + eβ = Offset(D)(β) Ia = boundary(β, N, Ip, false) Ib = boundary(β, N, Ip, true) Ja = Ia .+ eβ @@ -326,7 +322,7 @@ function apply_bc_p_pullback!(::PeriodicBC, φbar, β, t, setup; isright, kwargs isright && return φbar # We do both in one go for "left" (; dimension, N, Ip) = setup.grid D = dimension() - eβ = Offset{D}()(β) + eβ = Offset(D)(β) Ia = boundary(β, N, Ip, false) Ib = boundary(β, N, Ip, true) Ja = Ia .+ eβ @@ -372,7 +368,7 @@ function apply_bc_u!(bc::DirichletBC, u, β, t, setup; isright, dudt = false, kw ), D, ) - u[α][I] .= bcfunc.(α, xI..., t) + u[I, α] .= bcfunc.(α, xI..., t) end u end @@ -382,7 +378,7 @@ function apply_bc_u_pullback!(::DirichletBC, φbar, β, t, setup; isright, kwarg D = dimension() for α = 1:D I = boundary(β, N, Iu[α], isright) - φbar[α][I] .= 0 + φbar[I, α] .= 0 end φbar end @@ -390,7 +386,7 @@ end function apply_bc_p!(::DirichletBC, p, β, t, setup; isright, kwargs...) (; dimension, N, Ip) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) I = boundary(β, N, Ip, isright) J = isright ? I .- e(β) : I .+ e(β) @. p[I] = p[J] @@ -400,7 +396,7 @@ end function apply_bc_p_pullback!(::DirichletBC, φbar, β, t, setup; isright, kwargs...) (; dimension, N, Ip) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) I = boundary(β, N, Ip, isright) J = isright ? I .- e(β) : I .+ e(β) @. φbar[J] += φbar[I] @@ -434,12 +430,12 @@ end function apply_bc_u!(::SymmetricBC, u, β, t, setup; isright, kwargs...) (; dimension, Nu, Iu) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) for α = 1:D if α != β I = boundary(β, Nu[α], Iu[α], isright) J = isright ? I .- e(β) : I .+ e(β) - @. u[α][I] = u[α][J] + @. u[I, α] = u[J, α] end end u @@ -448,13 +444,13 @@ end function apply_bc_u_pullback!(::SymmetricBC, φbar, β, t, setup; isright, kwargs...) (; dimension, Nu, Iu) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) for α = 1:D if α != β I = boundary(β, Nu[α], Iu[α], isright) J = isright ? I .- e(β) : I .+ e(β) - @. φbar[α][J] += φbar[α][I] - @. φbar[α][I] = 0 + @. φbar[J, α] += φbar[I, α] + @. φbar[I, α] = 0 end end φbar @@ -463,7 +459,7 @@ end function apply_bc_p!(::SymmetricBC, p, β, t, setup; isright, kwargs...) (; dimension, N, Ip) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) I = boundary(β, N, Ip, isright) J = isright ? I .- e(β) : I .+ e(β) @. p[I] = p[J] @@ -473,7 +469,7 @@ end function apply_bc_p_pullback!(::SymmetricBC, φbar, β, t, setup; isright, kwargs...) (; dimension, N, Ip) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) I = boundary(β, N, Ip, isright) J = isright ? I .- e(β) : I .+ e(β) @. φbar[J] += φbar[I] @@ -490,11 +486,11 @@ apply_bc_temp_pullback!(bc::SymmetricBC, φbar, β, t, setup; isright, kwargs... function apply_bc_u!(bc::PressureBC, u, β, t, setup; isright, kwargs...) (; dimension, N, Iu) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) for α = 1:D I = boundary(β, N, Iu[α], isright) J = isright ? I .- e(β) : I .+ e(β) - @. u[α][I] = u[α][J] + @. u[I, α] = u[J, α] end u end @@ -502,12 +498,12 @@ end function apply_bc_u_pullback!(::PressureBC, φbar, β, t, setup; isright, kwargs...) (; dimension, N, Iu) = setup.grid D = dimension() - e = Offset{D}() + e = Offset(D) for α = 1:D I = boundary(β, N, Iu[α], isright) J = isright ? I .- e(β) : I .+ e(β) - @. φbar[α][J] += φbar[α][I] - @. φbar[α][I] = 0 + @. φbar[J, α] += φbar[I, α] + @. φbar[I, α] = 0 end φbar end diff --git a/src/eddyviscosity.jl b/src/eddyviscosity.jl index 0cbbdc527..8a0404e0f 100644 --- a/src/eddyviscosity.jl +++ b/src/eddyviscosity.jl @@ -1,45 +1,42 @@ function strain_natural!(S, u, setup) - (; grid, workgroupsize) = setup - (; dimension, Np, Ip, Δ, Δu) = grid - I0 = first(Ip) - I0 -= oneunit(I0) - strain_natural_kernel!(get_backend(u[1]), workgroupsize)(S, u, I0, Δ, Δu; ndrange = Np) + (; grid, backend, workgroupsize) = setup + (; Np, Ip, Δ, Δu) = grid + I0 = getoffset(Ip) + strain_natural_kernel!(backend, workgroupsize)(S, u, I0, Δ, Δu; ndrange = Np) S end -@kernel function strain_natural_kernel!(S, uu, I0::CartesianIndex{2}, Δ, Δu) +@kernel function strain_natural_kernel!(S, u, I0::CartesianIndex{2}, Δ, Δu) I = @index(Global, Cartesian) I = I + I0 - u, v = uu[1], uu[2] ex, ey = unit_cartesian_indices(2) Δux, Δuy = Δu[1][I[1]], Δ[2][I[2]] Δvx, Δvy = Δ[1][I[1]], Δu[2][I[2]] - ∂u∂x = (u[I] - u[I-ex]) / Δux - ∂u∂y = (u[I+ey] - u[I]) / Δuy - ∂v∂x = (v[I+ex] - v[I]) / Δvx - ∂v∂y = (v[I] - v[I-ey]) / Δvy + ∂u∂x = (u[I, 1] - u[I-ex, 1]) / Δux + ∂u∂y = (u[I+ey, 1] - u[I, 1]) / Δuy + ∂v∂x = (u[I+ex, 2] - u[I, 2]) / Δvx + ∂v∂y = (u[I, 2] - u[I-ey, 2]) / Δvy S.xx[I] = ∂u∂x S.yy[I] = ∂v∂y S.xy[I] = (∂u∂y + ∂v∂x) / 2 end -@kernel function strain_natural_kernel!(S, uu, I0::CartesianIndex{3}, Δ, Δu) +@kernel function strain_natural_kernel!(S, u, I0::CartesianIndex{3}, Δ, Δu) I = @index(Global, Cartesian) I = I + I0 - u, v, w = uu[1], uu[2], uu[3] ex, ey, ez = unit_cartesian_indices(3) Δux, Δuy, Δuz = Δu[1][I[1]], Δ[2][I[2]], Δ[3][I[3]] Δvx, Δvy, Δvz = Δ[1][I[1]], Δu[2][I[2]], Δ[3][I[3]] Δwx, Δwy, Δwz = Δ[1][I[1]], Δ[2][I[2]], Δu[3][I[3]] - ∂u∂x = (u[I] - u[I-ex]) / Δux - ∂u∂y = (u[I+ey] - u[I]) / Δuy - ∂u∂z = (u[I+ez] - u[I]) / Δuz - ∂v∂x = (v[I+ex] - v[I]) / Δvx - ∂v∂y = (v[I] - v[I-ey]) / Δvy - ∂v∂z = (v[I+ez] - v[I]) / Δvz - ∂w∂x = (w[I+ex] - w[I]) / Δwx - ∂w∂y = (w[I+ey] - w[I]) / Δwy - ∂w∂z = (w[I] - w[I-ez]) / Δwz + ∂u∂x = (u[I, 1] - u[I-ex, 1]) / Δux + ∂u∂y = (u[I+ey, 1] - u[I, 1]) / Δuy + ∂u∂z = (u[I+ez, 1] - u[I, 1]) / Δuz + ∂v∂x = (u[I+ex, 2] - u[I, 2]) / Δvx + ∂v∂y = (u[I, 2] - u[I-ey, 2]) / Δvy + ∂v∂z = (u[I+ez, 2] - u[I, 2]) / Δvz + ∂w∂x = (u[I+ex, 3] - u[I, 3]) / Δwx + ∂w∂y = (u[I+ey, 3] - u[I, 3]) / Δwy + ∂w∂z = (u[I, 3] - u[I-ez, 3]) / Δwz S.xx[I] = ∂u∂x S.yy[I] = ∂v∂y S.zz[I] = ∂w∂z @@ -49,18 +46,10 @@ end end function smagorinsky_viscosity!(visc, S, θ, setup) - (; grid, workgroupsize) = setup - (; dimension, Np, Ip, Δ) = grid - I0 = first(Ip) - I0 -= oneunit(I0) - smagorinsky_viscosity_kernel!(get_backend(visc), workgroupsize)( - visc, - S, - I0, - Δ, - θ; - ndrange = Np, - ) + (; grid, backend, workgroupsize) = setup + (; Np, Ip, Δ) = grid + I0 = getoffset(Ip) + smagorinsky_viscosity_kernel!(backend, workgroupsize)(visc, S, I0, Δ, θ; ndrange = Np) visc end @@ -90,18 +79,10 @@ end end function apply_eddy_viscosity!(σ, visc, setup) - (; grid, workgroupsize) = setup + (; grid, backend, workgroupsize) = setup (; Np, Ip, Δ, Δu) = grid - I0 = first(Ip) - I0 -= oneunit(I0) - apply_eddy_viscosity_kernel!(get_backend(visc), workgroupsize)( - σ, - visc, - I0, - Δ, - Δu; - ndrange = Np, - ) + I0 = getoffset(Ip) + apply_eddy_viscosity_kernel!(backend, workgroupsize)(σ, visc, I0, Δ, Δu; ndrange = Np) σ end @@ -109,8 +90,6 @@ end I = @index(Global, Cartesian) I = I + I0 ex, ey = unit_cartesian_indices(2) - Δpx, Δpy = Δ[1][I[1]], Δ[2][I[2]] - Δux, Δuy = Δu[1][I[1]], Δu[2][I[2]] # TODO: Add interpolation weights here visc_xy = (visc[I] + visc[I+ex] + visc[I+ey] + visc[I+ex+ey]) / 4 σ.xx[I] = 2 * visc[I] * σ.xx[I] @@ -122,8 +101,6 @@ end I = @index(Global, Cartesian) I = I + I0 ex, ey, ez = unit_cartesian_indices(3) - Δpx, Δpy, Δpz = Δ[1][I[1]], Δ[2][I[2]], Δ[3][I[3]] - Δux, Δuy, Δuz = Δu[1][I[1]], Δu[2][I[2]], Δu[3][I[3]] # TODO: Add interpolation weights here visc_xy = (visc[I] + visc[I+ex] + visc[I+ey] + visc[I+ex+ey]) / 4 visc_xz = (visc[I] + visc[I+ex] + visc[I+ez] + visc[I+ex+ez]) / 4 @@ -137,18 +114,10 @@ end end function divoftensor_natural!(c, σ, setup) - (; grid, workgroupsize) = setup + (; grid, backend, workgroupsize) = setup (; Np, Ip, Δ, Δu) = grid - I0 = first(Ip) - I0 -= oneunit(I0) - divoftensor_natural_kernel!(get_backend(c[1]), workgroupsize)( - c, - σ, - I0, - Δ, - Δu; - ndrange = Np, - ) + I0 = getoffset(Ip) + divoftensor_natural_kernel!(backend, workgroupsize)(c, σ, I0, Δ, Δu; ndrange = Np) c end @@ -162,8 +131,8 @@ end ∂σxy∂y = (σ.xy[I] - σ.xy[I-ey]) / Δpy ∂σyx∂x = (σ.xy[I] - σ.xy[I-ex]) / Δpx ∂σyy∂y = (σ.yy[I+ey] - σ.yy[I]) / Δuy - c[1][I] = ∂σxx∂x + ∂σxy∂y - c[2][I] = ∂σyx∂x + ∂σyy∂y + c[I, 1] = ∂σxx∂x + ∂σxy∂y + c[I, 2] = ∂σyx∂x + ∂σyy∂y end @kernel function divoftensor_natural_kernel!(c, σ, I0::CartesianIndex{3}, Δ, Δu) @@ -181,9 +150,9 @@ end ∂σzx∂x = (σ.xz[I] - σ.xz[I-ex]) / Δpx ∂σzy∂y = (σ.yz[I] - σ.yz[I-ey]) / Δpy ∂σzz∂z = (σ.zz[I+ez] - σ.zz[I]) / Δuz - c[1][I] = ∂σxx∂x + ∂σxy∂y + ∂σxz∂z - c[2][I] = ∂σyx∂x + ∂σyy∂y + ∂σyz∂z - c[3][I] = ∂σzx∂x + ∂σzy∂y + ∂σzz∂z + c[I, 1] = ∂σxx∂x + ∂σxy∂y + ∂σxz∂z + c[I, 2] = ∂σyx∂x + ∂σyy∂y + ∂σyz∂z + c[I, 3] = ∂σzx∂x + ∂σzy∂y + ∂σzz∂z end function smagorinsky_closure_natural(setup) diff --git a/src/initializers.jl b/src/initializers.jl index 0c437b4dc..d6e6cc2e7 100644 --- a/src/initializers.jl +++ b/src/initializers.jl @@ -3,7 +3,7 @@ scalarfield(setup) = fill!(similar(setup.grid.x[1], setup.grid.N), 0) "Create empty vector field." vectorfield(setup) = - ntuple(α -> fill!(similar(setup.grid.x[1], setup.grid.N), 0), setup.grid.dimension()) + fill!(similar(setup.grid.x[1], setup.grid.N..., setup.grid.dimension()), 0) """ Create divergence free velocity field `u` with boundary conditions at time `t`. @@ -31,7 +31,7 @@ function velocityfield( β -> reshape(xu[α][β][Iu[α].indices[β]], ntuple(Returns(1), β - 1)..., :), D, ) - u[α][Iu[α]] .= ufunc.(α, xin...) + u[Iu[α], α] .= ufunc.(α, xin...) end # Make velocity field divergence free @@ -177,6 +177,7 @@ function create_spectrum(; setup, kp, rng = Random.default_rng()) # end a .* eα end + stack(uhat) end """ @@ -201,16 +202,16 @@ function random_field( all(==((PeriodicBC(), PeriodicBC())), boundary_conditions), "Random field requires periodic boundary conditions." ) - @assert all(Δ -> all(≈(0), diff(Δ)), Δ) "Random field requires uniform grid spacing." + @assert all(Δ -> all(≈(Δ[1]), Δ), Δ) "Random field requires uniform grid spacing." @assert all(iseven, N) "Random field requires even number of volumes." # Create random velocity field uhat = create_spectrum(; setup, kp, rng) - u = ifft.(uhat) - u = map(u -> A .* real.(u), u) + u = ifft(uhat, 1:D) + u = @. A * real(u) # Add ghost volumes (one on each side for periodic) - u = pad_circular.(u, 1; dims = 1:D) + u = pad_circular(u, 1; dims = 1:D) # # Interpolate to staggered grid # interpolate_p_u!(u, setup) diff --git a/src/matrices.jl b/src/matrices.jl index a5405f767..437d8ea9d 100644 --- a/src/matrices.jl +++ b/src/matrices.jl @@ -5,7 +5,7 @@ function laplacian_mat(setup) backend = get_backend(x[1]) T = eltype(x[1]) D = dimension() - e = Offset{D}() + e = Offset(D) Ia = first(Ip) Ib = last(Ip) I = similar(x[1], CartesianIndex{D}, 0) diff --git a/src/operators.jl b/src/operators.jl index 4f16779bd..01e54ef41 100644 --- a/src/operators.jl +++ b/src/operators.jl @@ -99,32 +99,27 @@ divergence(u, setup) = divergence!(scalarfield(setup), u, setup) ChainRulesCore.rrule(::typeof(divergence), u, setup) = ( divergence(u, setup), - φ -> ( - NoTangent(), - Tangent{typeof(u)}(divergence_adjoint!(vectorfield(setup), φ, setup)...), - NoTangent(), - ), + φ -> (NoTangent(), divergence_adjoint!(vectorfield(setup), φ, setup), NoTangent()), ) "Compute divergence of velocity field (in-place version)." function divergence!(div, u, setup) (; grid, backend, workgroupsize) = setup (; Δ, Ip, Np) = grid - D = length(u) + D = length(Δ) e = Offset(D) I0 = getoffset(Ip) kernel = divergence_kernel!(backend, workgroupsize) kernel(div, u, Δ, e, I0; ndrange = Np) - # kernel(div, u, Δ, e, Val(1:D), true, I0; ndrange = Np) div end -@kernel function divergence_kernel!(div, u, Δ, e, I0) +@kernel inbounds = true function divergence_kernel!(div, u, Δ, e, I0) I = @index(Global, Cartesian) I = I + I0 d = zero(eltype(div)) - for α in eachindex(u) - d += (u[α][I] - u[α][I-e(α)]) / Δ[α][I[α]] + for α in eachindex(Δ) + d += (u[I, α] - u[I-e(α), α]) / Δ[α][I[α]] end div[I] = d end @@ -132,7 +127,7 @@ end function divergence_adjoint!(u, φ, setup) (; grid, backend, workgroupsize) = setup (; Δ, N, Ip) = grid - D = length(u) + D = length(Δ) e = Offset(D) divergence_adjoint_kernel!(backend, workgroupsize)(u, φ, Δ, Ip, e; ndrange = N) u @@ -141,10 +136,10 @@ end @kernel function divergence_adjoint_kernel!(u, φ, Δ, Ip, e) I = @index(Global, Cartesian) for α in eachindex(Δ) - adjoint = zero(eltype(u[1])) + adjoint = zero(eltype(u)) I ∈ Ip && (adjoint += φ[I] / Δ[α][I[α]]) I + e(α) ∈ Ip && (adjoint -= φ[I+e(α)] / Δ[α][I[α]+1]) - u[α][I] += adjoint + u[I, α] += adjoint end end @@ -155,7 +150,7 @@ ChainRulesCore.rrule(::typeof(pressuregradient), p, setup) = ( pressuregradient(p, setup), φ -> ( NoTangent(), - pressuregradient_adjoint!(scalarfield(setup), (φ...,), setup), + pressuregradient_adjoint!(scalarfield(setup), φ, setup), NoTangent(), ), ) @@ -166,7 +161,6 @@ function pressuregradient!(G, p, setup) (; dimension, Δu, N, Iu) = grid D = dimension() e = Offset(D) - D = dimension() kernel = pressuregradient_kernel!(backend, workgroupsize) I0 = oneunit(CartesianIndex{D}) kernel(G, p, Δu, Iu, e, Val(1:D), I0; ndrange = N .- 2) @@ -178,7 +172,7 @@ end I = I0 + I @unroll for α in getval(valdims) if I ∈ Iu[α] - G[α][I] = (p[I+e(α)] - p[I]) / Δu[α][I[α]] + G[I, α] = (p[I+e(α)] - p[I]) / Δu[α][I[α]] end end end @@ -197,8 +191,8 @@ end I = @index(Global, Cartesian) adjoint = zero(eltype(p)) @unroll for α in getval(valdims) - I - e(α) ∈ Iu[α] && (adjoint += φ[α][I-e(α)] / Δu[α][I[α]-1]) - I ∈ Iu[α] && (adjoint -= φ[α][I] / Δu[α][I[α]]) + I - e(α) ∈ Iu[α] && (adjoint += φ[I-e(α), α] / Δu[α][I[α]-1]) + I ∈ Iu[α] && (adjoint -= φ[I, α] / Δu[α][I[α]]) end p[I] += adjoint end @@ -227,7 +221,7 @@ end I = @index(Global, Cartesian) I = I0 + I @unroll for α in getval(valdims) - u[α][I] -= (p[I+e(α)] - p[I]) / Δu[α][I[α]] + u[I, α] -= (p[I+e(α)] - p[I]) / Δu[α][I[α]] end end @@ -315,8 +309,7 @@ function laplacian!(L, p, setup) # All volumes have a left velocity except the first one # Start at second volume ndrange = Np - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) # lap!(backend, workgroupsize)(L, p, I0; ndrange) L .= 0 for α = 1:D @@ -326,15 +319,11 @@ function laplacian!(L, p, setup) end "Compute convective term (differentiable version)." -convection(u, setup) = convection!(zero.(u), u, setup) +convection(u, setup) = convection!(zero(u), u, setup) ChainRulesCore.rrule(::typeof(convection), u, setup) = ( convection(u, setup), - φ -> ( - NoTangent(), - Tangent{typeof(u)}(convection_adjoint!(zero.(u), (φ...,), u, setup)...), - NoTangent(), - ), + φ -> (NoTangent(), convection_adjoint!(zero(u), φ, u, setup), NoTangent()), ) """ @@ -357,7 +346,7 @@ end I = @index(Global, Cartesian) I = I + I0 @unroll for α in dims - f = F[α][I] + f = F[I, α] if I ∈ Iu[α] @unroll for β in dims Δuαβ = α == β ? Δu[β] : Δ[β] @@ -367,16 +356,16 @@ end # In matrix version, uses # 1*u[α][I-e(β)] + 0*u[α][I] # instead of 1/2 when u[α][I-e(β)] is at Dirichlet boundary. - uαβ1 = (u[α][I-e(β)] + u[α][I]) / 2 - uαβ2 = (u[α][I] + u[α][I+e(β)]) / 2 + uαβ1 = (u[I-e(β), α] + u[I, α]) / 2 + uαβ2 = (u[I, α] + u[I+e(β), α]) / 2 uβα1 = - A[β][α][2][I[α]-(α==β)] * u[β][I-e(β)] + - A[β][α][1][I[α]+(α!=β)] * u[β][I-e(β)+e(α)] - uβα2 = A[β][α][2][I[α]] * u[β][I] + A[β][α][1][I[α]+1] * u[β][I+e(α)] + A[β][α][2][I[α]-(α==β)] * u[I-e(β), β] + + A[β][α][1][I[α]+(α!=β)] * u[I-e(β)+e(α), β] + uβα2 = A[β][α][2][I[α]] * u[I, β] + A[β][α][1][I[α]+1] * u[I+e(α), β] f -= (uαβ2 * uβα2 - uαβ1 * uβα1) / Δuαβ[I[β]] end end - F[α][I] = f + F[I, α] = f end end @@ -385,7 +374,7 @@ function convection_adjoint!(ubar, φbar, u, setup) (; dimension, Δ, Δu, N, Iu, A) = grid D = dimension() e = Offset(D) - T = eltype(u[1]) + T = eltype(u) h = T(1) / 2 kernel = convection_adjoint_kernel!(backend, workgroupsize) kernel(ubar, φbar, u, Δ, Δu, Iu, A, h, e, Val(1:D); ndrange = N) @@ -396,7 +385,7 @@ end dims = getval(valdims) J = @index(Global, Cartesian) @unroll for γ in dims - adjoint = zero(eltype(u[1])) + adjoint = zero(eltype(u)) @unroll for α in dims @unroll for β in dims Δuαβ = α == β ? Δu[β] : Δ[β] @@ -407,36 +396,36 @@ end I = J if α == γ && I in Iu[α] uαβ2 = h - uβα2 = Aβα2[I[α]] * u[β][I] + Aβα1[I[α]+1] * u[β][I+e(α)] + uβα2 = Aβα2[I[α]] * u[I, β] + Aβα1[I[α]+1] * u[I+e(α), β] dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 2 I = J - e(β) if α == γ && I in Iu[α] uαβ2 = h - uβα2 = Aβα2[I[α]] * u[β][I] + Aβα1[I[α]+1] * u[β][I+e(α)] + uβα2 = Aβα2[I[α]] * u[I, β] + Aβα1[I[α]+1] * u[I+e(α), β] dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 3 I = J if β == γ && I in Iu[α] - uαβ2 = h * u[α][I] + h * u[α][I+e(β)] + uαβ2 = h * u[I, α] + h * u[I+e(β), α] uβα2 = Aβα2[I[α]] dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 4 I = J - e(α) if β == γ && I in Iu[α] - uαβ2 = h * u[α][I] + h * u[α][I+e(β)] + uαβ2 = h * u[I, α] + h * u[I+e(β), α] uβα2 = Aβα1[I[α]+1] dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 5 @@ -444,10 +433,10 @@ end if α == γ && I in Iu[α] uαβ1 = h uβα1 = - Aβα2[I[α]-(α==β)] * u[β][I-e(β)] + - Aβα1[I[α]+(α!=β)] * u[β][I-e(β)+e(α)] + Aβα2[I[α]-(α==β)] * u[I-e(β), β] + + Aβα1[I[α]+(α!=β)] * u[I-e(β)+e(α), β] dφdu = uαβ1 * uβα1 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 6 @@ -455,32 +444,32 @@ end if α == γ && I in Iu[α] uαβ1 = h uβα1 = - Aβα2[I[α]-(α==β)] * u[β][I-e(β)] + - Aβα1[I[α]+(α!=β)] * u[β][I-e(β)+e(α)] + Aβα2[I[α]-(α==β)] * u[I-e(β), β] + + Aβα1[I[α]+(α!=β)] * u[I-e(β)+e(α), β] dφdu = uαβ1 * uβα1 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 7 I = J + e(β) if β == γ && I in Iu[α] - uαβ1 = h * u[α][I-e(β)] + h * u[α][I] + uαβ1 = h * u[I-e(β), α] + h * u[I, α] uβα1 = Aβα2[I[α]-(α==β)] dφdu = uαβ1 * uβα1 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end # 8 I = J + e(β) - e(α) if β == γ && I in Iu[α] - uαβ1 = h * u[α][I-e(β)] + h * u[α][I] + uαβ1 = h * u[I-e(β), α] + h * u[I, α] uβα1 = Aβα1[I[α]+(α!=β)] dφdu = uαβ1 * uβα1 / Δuαβ[I[β]] - adjoint += φbar[α][I] * dφdu + adjoint += φbar[I, α] * dφdu end end end - ubar[γ][J] += adjoint + ubar[J, γ] += adjoint end end @@ -489,11 +478,7 @@ diffusion(u, setup) = diffusion!(zero.(u), u, setup) ChainRulesCore.rrule(::typeof(diffusion), u, setup) = ( diffusion(u, setup), - φ -> ( - NoTangent(), - Tangent{typeof(u)}(diffusion_adjoint!(zero.(u), (φ...,), setup)...), - NoTangent(), - ), + φ -> (NoTangent(), diffusion_adjoint!(zero(u), φ, setup), NoTangent()), ) """ @@ -517,18 +502,18 @@ end I = @index(Global, Cartesian) I = I + I0 @unroll for α in dims - f = F[α][I] + f = F[I, α] if I ∈ Iu[α] @unroll for β in dims Δuαβ = α == β ? Δu[β] : Δ[β] Δa = β == α ? Δ[β][I[β]] : Δu[β][I[β]-1] Δb = β == α ? Δ[β][I[β]+1] : Δu[β][I[β]] - ∂a = (u[α][I] - u[α][I-e(β)]) / Δa - ∂b = (u[α][I+e(β)] - u[α][I]) / Δb + ∂a = (u[I, α] - u[I-e(β), α]) / Δa + ∂b = (u[I+e(β), α] - u[I, α]) / Δb f += visc * (∂b - ∂a) / Δuαβ[I[β]] end end - F[α][I] = f + F[I, α] = f end end @@ -550,28 +535,28 @@ end val = zero(eltype(u[1])) @unroll for β in dims Δuαβ = α == β ? Δu[β] : Δ[β] - # F[α][I] += visc * u[α][I+e(β)] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) - # F[α][I] -= visc * u[α][I] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) - # F[α][I] -= visc * u[α][I] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) - # F[α][I] += visc * u[α][I-e(β)] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) + # F[α][I] += visc * u[I+e(β), α] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) + # F[α][I] -= visc * u[I, α] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) + # F[α][I] -= visc * u[I, α] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) + # F[α][I] += visc * u[I-e(β), α] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) if I - e(β) ∈ Iu[α] val += - visc * φ[α][I-e(β)] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / + visc * φ[I-e(β), α] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / Δuαβ[I[β]-1] end if I ∈ Iu[α] - val -= visc * φ[α][I] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / Δuαβ[I[β]] + val -= visc * φ[I, α] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / Δuαβ[I[β]] end if I ∈ Iu[α] - val -= visc * φ[α][I] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / Δuαβ[I[β]] + val -= visc * φ[I, α] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / Δuαβ[I[β]] end if I + e(β) ∈ Iu[α] val += - visc * φ[α][I+e(β)] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / + visc * φ[I+e(β), α] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / Δuαβ[I[β]+1] end end - u[α][I] += val + u[I, α] += val end end @@ -582,7 +567,7 @@ end # convection(u, setup), # φ -> ( # NoTangent(), -# Tangent{typeof(u)}(convectiondiffusion_adjoint!(vectorfield(setup), φ, setup)...), +# convectiondiffusion_adjoint!(vectorfield(setup), φ, setup), # NoTangent(), # ), # ) @@ -596,8 +581,7 @@ function convectiondiffusion!(F, u, setup) (; dimension, Δ, Δu, N, A, Iu) = grid D = dimension() e = Offset(D) - @assert all(==(N), size.(F)) - @assert all(==(N), size.(u)) + @assert size(u) == size(F) == (N..., D) visc = 1 / Re I0 = oneunit(CartesianIndex{D}) kernel = convection_diffusion_kernel!(backend, workgroupsize) @@ -621,24 +605,24 @@ end I = I + I0 dims = getval(valdims) @unroll for α in dims - f = F[α][I] + f = F[I, α] if I ∈ Iu[α] @unroll for β in dims Δuαβ = α == β ? Δu[β] : Δ[β] - uαβ1 = (u[α][I-e(β)] + u[α][I]) / 2 - uαβ2 = (u[α][I] + u[α][I+e(β)]) / 2 + uαβ1 = (u[I-e(β), α] + u[I, α]) / 2 + uαβ2 = (u[I, α] + u[I+e(β), α]) / 2 uβα1 = - A[β][α][2][I[α]-(α==β)] * u[β][I-e(β)] + - A[β][α][1][I[α]+(α!=β)] * u[β][I-e(β)+e(α)] - uβα2 = A[β][α][2][I[α]] * u[β][I] + A[β][α][1][I[α]+1] * u[β][I+e(α)] + A[β][α][2][I[α]-(α==β)] * u[I-e(β), β] + + A[β][α][1][I[α]+(α!=β)] * u[I-e(β)+e(α), β] + uβα2 = A[β][α][2][I[α]] * u[I, β] + A[β][α][1][I[α]+1] * u[I+e(α), β] uαuβ1 = uαβ1 * uβα1 uαuβ2 = uαβ2 * uβα2 - ∂βuα1 = (u[α][I] - u[α][I-e(β)]) / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) - ∂βuα2 = (u[α][I+e(β)] - u[α][I]) / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) + ∂βuα1 = (u[I, α] - u[I-e(β), α]) / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) + ∂βuα2 = (u[I+e(β), α] - u[I, α]) / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) f += (visc * (∂βuα2 - ∂βuα1) - (uαuβ2 - uαuβ1)) / Δuαβ[I[β]] end end - F[α][I] = f + F[I, α] = f end end @@ -647,7 +631,7 @@ Compute convection-diffusion term for the temperature equation. (differentiable version). """ convection_diffusion_temp(u, temp, setup) = - convection_diffusion_temp!(zero.(temp), u, temp, setup) + convection_diffusion_temp!(zero(temp), u, temp, setup) function ChainRulesCore.rrule(::typeof(convection_diffusion_temp), u, temp, setup) conv = convection_diffusion_temp(u, temp, setup) @@ -679,60 +663,15 @@ end @unroll for β in getval(valdims) ∂T∂x1 = (temp[I] - temp[I-e(β)]) / Δu[β][I[β]-1] ∂T∂x2 = (temp[I+e(β)] - temp[I]) / Δu[β][I[β]] - uT1 = u[β][I-e(β)] * avg(temp, Δ, I - e(β), β) - uT2 = u[β][I] * avg(temp, Δ, I, β) + uT1 = u[I-e(β), β] * avg(temp, Δ, I - e(β), β) + uT2 = u[I, β] * avg(temp, Δ, I, β) cI += (-(uT2 - uT1) + α4 * (∂T∂x2 - ∂T∂x1)) / Δ[β][I[β]] end c[I] += cI end -# function convection_diffusion_temp_pullback(u, temp, setup) -# (; grid, backend, workgroupsize, temperature) = setup -# (; dimension, Δ, Δu, Np, Ip) = grid -# (; α4) = temperature -# D = dimension() -# e = Offset(D) -# @kernel function pullback_u!(du, φ, u, temp, ::Val{βrange}, I0) where {βrange} -# I = @index(Global, Cartesian) -# I = I + I0 -# cI = zero(eltype(c)) -# @unroll for β in βrange -# ∂T∂x1 = (temp[I] - temp[I-e(β)]) / Δu[β][I[β]-1] -# ∂T∂x2 = (temp[I+e(β)] - temp[I]) / Δu[β][I[β]] -# uT1 = u[β][I-e(β)] * avg(temp, Δ, I - e(β), β) -# uT2 = u[β][I] * avg(temp, Δ, I, β) -# cI += (-(uT2 - uT1) + α4 * (∂T∂x2 - ∂T∂x1)) / Δ[β][I[β]] -# end -# c[I] = cI -# end -# I0 = first(Ip) -# I0 -= oneunit(I0) -# conv!(backend, workgroupsize)(c, u, temp, Val(1:D), I0; ndrange = Np) -# function pullback(φ) -# (NoTangent(), du, dtemp, NoTangent()) -# end -# end - -# function dissipation!(c, u, setup) -# (; grid, workgroupsize, temperature) = setup -# (; dimension, Δ, Np, Ip) = grid -# D = dimension() -# e = Offset(D) -# @inline ∂2(u, α, β, I) = ((u[α][I+e(β)] - u[α][I]) / Δ[β][I])^2 / 2 -# @inline Φ(u, α, β, I) = -∂2(u, α, β, I) - ∂2(u, α, β, I+e(β)) -# @kernel function diss!(d, u, ::Val{βrange}, I0) where {βrange} -# I = @index(Global, Cartesian) -# I = I + I0 -# cI = zero(eltype(c)) -# @unroll for β in βrange -# cI += Φ(u, β, β, I) / Δ[β][I[β]] -# end -# c[I] += cI -# end -# end - "Compute dissipation term for the temperature equation (differentiable version)." -dissipation(u, setup) = dissipation!(zero(u[1]), zero.(u), u, setup) +dissipation(u, setup) = dissipation!(scalarfield(setup), vectorfield(setup), u, setup) function ChainRulesCore.rrule(::typeof(dissipation), u, setup) (; grid, backend, workgroupsize, Re, temperature) = setup @@ -741,38 +680,37 @@ function ChainRulesCore.rrule(::typeof(dissipation), u, setup) D = dimension() e = Offset(D) d, d_pb = ChainRulesCore.rrule(diffusion, u, setup) - φ = dissipation!(zero(u[1]), d, u, setup) - @kernel function ∂φ!(ubar, dbar, φbar, d, u, ::Val{βrange}) where {βrange} + φ = dissipation!(scalarfield(setup), d, u, setup) + @kernel function ∂φ!(ubar, dbar, φbar, d, u, valdims) J = @index(Global, Cartesian) - @unroll for β in βrange + @unroll for β in getval(valdims) # Compute ubar - a = zero(eltype(u[1])) + a = zero(eltype(u)) # 1 I = J + e(β) - I ∈ Ip && (a += Re * α1 / γ * d[β][I-e(β)] / 2) + I ∈ Ip && (a += Re * α1 / γ * d[I-e(β), β] / 2) # 2 I = J - I ∈ Ip && (a += Re * α1 / γ * d[β][I] / 2) - ubar[β][J] += a + I ∈ Ip && (a += Re * α1 / γ * d[I, β] / 2) + ubar[J, β] += a # Compute dbar - b = zero(eltype(u[1])) + b = zero(eltype(u)) # 1 I = J + e(β) - I ∈ Ip && (b += Re * α1 / γ * u[β][I-e(β)] / 2) + I ∈ Ip && (b += Re * α1 / γ * u[I-e(β), β] / 2) # 2 I = J - I ∈ Ip && (b += Re * α1 / γ * u[β][I] / 2) - dbar[β][J] += b + I ∈ Ip && (b += Re * α1 / γ * u[I, β] / 2) + dbar[J, β] += b end end function dissipation_pullback(φbar) # Dφ/Du = ∂φ(u, d)/∂u + ∂φ(u, d)/∂d ⋅ ∂d(u)/∂u - dbar = zero.(u) - ubar = zero.(u) + dbar = zero(u) + ubar = zero(u) ∂φ!(backend, workgroupsize)(ubar, dbar, φbar, d, u, Val(1:D); ndrange = N) diffusion_adjoint!(ubar, dbar, setup) - ubar = Tangent{typeof(u)}(ubar...) (NoTangent(), ubar, NoTangent()) end φ, dissipation_pullback @@ -788,14 +726,14 @@ function dissipation!(diss, diff, u, setup) (; α1, γ) = temperature D = dimension() e = Offset(D) - fill!.(diff, 0) + fill!(diff, 0) diffusion!(diff, u, setup) - @kernel function interpolate!(diss, diff, u, I0, ::Val{βrange}) where {βrange} + @kernel function interpolate!(diss, diff, u, I0, valdims) I = @index(Global, Cartesian) I += I0 d = zero(eltype(diss)) - @unroll for β in βrange - d += Re * α1 / γ * (u[β][I-e(β)] * diff[β][I-e(β)] + u[β][I] * diff[β][I]) / 2 + @unroll for β in getval(valdims) + d += Re * α1 / γ * (u[I-e(β), β] * diff[I-e(β), β] + u[I, β] * diff[I, β]) / 2 end diss[I] += d end @@ -810,7 +748,7 @@ Compute dissipation term ``2 \\nu \\langle S_{i j} S_{i j} \\rangle`` from strain-rate tensor (differentiable version). """ -dissipation_from_strain(u, setup) = dissipation_from_strain!(zero(u[1]), u, setup) +dissipation_from_strain(u, setup) = dissipation_from_strain!(scalarfield(setup), u, setup) ChainRulesCore.rrule(::typeof(dissipation_from_strain), u, setup) = (dissipation_from_strain(u, setup), φ -> error("Not yet implemented")) @@ -836,12 +774,12 @@ end "Compute body force (differentiable version)." function applybodyforce(u, t, setup) (; grid, bodyforce, issteadybodyforce) = setup - (; dimension) = grid + (; dimension, x) = grid D = dimension() if issteadybodyforce bodyforce else - map(α -> bodyforce.(α, x[α]..., t), 1:D) + stack(map(α -> bodyforce.(α, x[α]..., t), 1:D)) end end @@ -859,9 +797,9 @@ function applybodyforce!(F, u, t, setup) (; grid, bodyforce, issteadybodyforce) = setup (; dimension, Iu, xu) = grid D = dimension() - for α = 1:D + for (α, Fα) in enumerate(eachslice(F; dims = D + 1)) if issteadybodyforce - F[α] .+= bodyforce[α] + F .+= bodyforce else # xin = ntuple( # β -> reshape(xu[α][β][Iu[α].indices[β]], ntuple(Returns(1), β - 1)..., :), @@ -869,15 +807,14 @@ function applybodyforce!(F, u, t, setup) # ) # @. F[α][Iu[α]] += bodyforce(α, xin..., t) xin = ntuple(β -> reshape(xu[α][β], ntuple(Returns(1), β - 1)..., :), D) - F[α] .+= bodyforce.(α, xin..., t) + Fα .+= bodyforce.(α, xin..., t) end end F end "Compute gravity term (differentiable version)." -gravity(temp, setup) = - gravity!(ntuple(α -> zero(temp), setup.grid.dimension()), temp, setup) +gravity(temp, setup) = gravity!(vectorfield(setup), temp, setup) function ChainRulesCore.rrule(::typeof(gravity), temp, setup) (; grid, backend, workgroupsize, temperature) = setup @@ -888,15 +825,16 @@ function ChainRulesCore.rrule(::typeof(gravity), temp, setup) e = Offset(D) g = gravity(temp, setup) function gravity_pullback(φ) - @kernel function g!(tempbar, φbar, ::Val{α}) where {α} + @kernel function g!(tempbar, φbar, valα) + α = getval(valα) J = @index(Global, Cartesian) t = zero(eltype(tempbar)) # 1 I = J - I ∈ Iu[α] && (t += α2 * Δ[α][I[α]+1] * φbar[α][I] / (Δ[α][I[α]] + Δ[α][I[α]+1])) + I ∈ Iu[α] && (t += α2 * Δ[α][I[α]+1] * φbar[I, α] / (Δ[α][I[α]] + Δ[α][I[α]+1])) # 2 I = J - e(α) - I ∈ Iu[α] && (t += α2 * Δ[α][I[α]] * φbar[α][I] / (Δ[α][I[α]] + Δ[α][I[α]+1])) + I ∈ Iu[α] && (t += α2 * Δ[α][I[α]] * φbar[I, α] / (Δ[α][I[α]] + Δ[α][I[α]+1])) tempbar[J] = t end tempbar = zero(temp) @@ -919,7 +857,7 @@ function gravity!(F, temp, setup) @kernel function g!(F, temp, ::Val{gdir}, I0) where {gdir} I = @index(Global, Cartesian) I = I + I0 - F[gdir][I] += α2 * avg(temp, Δ, I, gdir) + F[I, gdir] += α2 * avg(temp, Δ, I, gdir) end I0 = first(Iu[gdir]) I0 -= oneunit(I0) @@ -967,7 +905,7 @@ function momentum!(F, u, temp, t, setup) (; grid, closure_model, bodyforce, temperature) = setup (; dimension) = grid D = dimension() - fill!.(F, 0) + fill!(F, 0) # diffusion!(F, u, setup) # convection!(F, u, setup) convectiondiffusion!(F, u, setup) @@ -985,8 +923,11 @@ end # (tupleadd(u...), φ -> (NoTangent(), map(u -> φ, u)...)) "Compute vorticity field (differentiable version)." -vorticity(u, setup) = - vorticity!(length(u) == 2 ? scalarfield(setup) : vectorfield(setup), u, setup) +vorticity(u, setup) = vorticity!( + setup.grid.dimension() == 2 ? scalarfield(setup) : vectorfield(setup), + u, + setup, +) "Compute vorticity field (in-place version)." vorticity!(ω, u, setup) = vorticity!(setup.grid.dimension, ω, u, setup) @@ -1000,7 +941,7 @@ function vorticity!(::Dimension{2}, ω, u, setup) @kernel function ω!(ω, u) I = @index(Global, Cartesian) ω[I] = - (u[2][I+e(1)] - u[2][I]) / Δu[1][I[1]] - (u[1][I+e(2)] - u[1][I]) / Δu[2][I[2]] + (u[I+e(1), 2] - u[I, 2]) / Δu[1][I[1]] - (u[I+e(2), 1] - u[I, 1]) / Δu[2][I[2]] end ω!(backend, workgroupsize)(ω, u; ndrange = N .- 1) ω @@ -1017,31 +958,31 @@ function vorticity!(::Dimension{3}, ω, u, setup) for (α, α₊, α₋) in ((1, 2, 3), (2, 3, 1), (3, 1, 2)) # α₊ = mod1(α + 1, D) # α₋ = mod1(α - 1, D) - ω[α][I] = - (u[α₋][I+e(α₊)] - u[α₋][I]) / Δu[α₊][I[α₊]] - - (u[α₊][I+e(α₋)] - u[α₊][I]) / Δu[α₋][I[α₋]] + ω[I, α] = + (u[I+e(α₊), α₋] - u[I, α₋]) / Δu[α₊][I[α₊]] - + (u[I+e(α₋), α₊] - u[I, α₊]) / Δu[α₋][I[α₋]] end end ω!(backend, workgroupsize)(ω, u; ndrange = N .- 1) ω end -@inline ∂x(uα, I::CartesianIndex{D}, α, β, Δβ, Δuβ; e = Offset(D)) where {D} = - α == β ? (uα[I] - uα[I-e(β)]) / Δβ[I[β]] : +@inline ∂x(u, I::CartesianIndex{D}, α, β, Δβ, Δuβ; e = Offset(D)) where {D} = + α == β ? (u[I, α] - u[I-e(β), α]) / Δβ[I[β]] : ( - (uα[I+e(β)] - uα[I]) / Δuβ[I[β]] + - (uα[I-e(α)+e(β)] - uα[I-e(α)]) / Δuβ[I[β]] + - (uα[I] - uα[I-e(β)]) / Δuβ[I[β]-1] + - (uα[I-e(α)] - uα[I-e(α)-e(β)]) / Δuβ[I[β]-1] + (u[I+e(β), α] - u[I, α]) / Δuβ[I[β]] + + (u[I-e(α)+e(β), α] - u[I-e(α), α]) / Δuβ[I[β]] + + (u[I, α] - u[I-e(β), α]) / Δuβ[I[β]-1] + + (u[I-e(α), α] - u[I-e(α)-e(β), α]) / Δuβ[I[β]-1] ) / 4 @inline ∇(u, I::CartesianIndex{2}, Δ, Δu) = - @SMatrix [∂x(u[α], I, α, β, Δ[β], Δu[β]) for α = 1:2, β = 1:2] + @SMatrix [∂x(u, I, α, β, Δ[β], Δu[β]) for α = 1:2, β = 1:2] @inline ∇(u, I::CartesianIndex{3}, Δ, Δu) = - @SMatrix [∂x(u[α], I, α, β, Δ[β], Δu[β]) for α = 1:3, β = 1:3] + @SMatrix [∂x(u, I, α, β, Δ[β], Δu[β]) for α = 1:3, β = 1:3] @inline idtensor(u, I::CartesianIndex{2}) = - @SMatrix [(α == β) * oneunit(eltype(u[1])) for α = 1:2, β = 1:2] + @SMatrix [(α == β) * oneunit(eltype(u)) for α = 1:2, β = 1:2] @inline idtensor(u, I::CartesianIndex{3}) = - @SMatrix [(α == β) * oneunit(eltype(u[1])) for α = 1:3, β = 1:3] + @SMatrix [(α == β) * oneunit(eltype(u)) for α = 1:3, β = 1:3] @inline function strain(u, I, Δ, Δu) ∇u = ∇(u, I, Δ, Δu) (∇u + ∇u') / 2 @@ -1065,8 +1006,7 @@ function smagtensor!(σ, u, θ, setup) eddyvisc = θ^2 * d^2 * sqrt(2 * sum(S .* S)) σ[I] = 2 * eddyvisc * S end - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) σ!(backend, workgroupsize)(σ, u, I0; ndrange = Np) σ end @@ -1084,7 +1024,7 @@ function divoftensor!(s, σ, setup) @kernel function s!(s, σ, ::Val{α}, ::Val{βrange}, I0) where {α,βrange} I = @index(Global, Cartesian) I = I + I0 - s[α][I] = zero(eltype(s[1])) + s[I, α] = zero(eltype(s[1])) # for β = 1:D @unroll for β in βrange Δuαβ = α == β ? Δu[β] : Δ[β] @@ -1108,12 +1048,11 @@ function divoftensor!(s, σ, setup) σ[I+e(α)][α, β] ) / 4 end - s[α][I] += (σαβ2 - σαβ1) / Δuαβ[I[β]] + s[I, α] += (σαβ2 - σαβ1) / Δuαβ[I[β]] end end for α = 1:D - I0 = first(Iu[α]) - I0 -= oneunit(I0) + I0 = getoffset(Iu[α]) s!(backend, workgroupsize)(s, σ, Val(α), Val(1:D), I0; ndrange = Nu[α]) end s @@ -1139,11 +1078,11 @@ end "Compute symmetry tensor basis (differentiable version)." function tensorbasis(u, setup) - T = eltype(u[1]) + T = eltype(u) D = setup.grid.dimension() tensorbasis!( - ntuple(α -> similar(u[1], SMatrix{D,D,T,D * D}, setup.grid.N), D == 2 ? 3 : 11), - ntuple(α -> similar(u[1], setup.grid.N), D == 2 ? 2 : 5), + ntuple(α -> similar(u, SMatrix{D,D,T,D * D}, setup.grid.N), D == 2 ? 3 : 11), + ntuple(α -> similar(u, setup.grid.N), D == 2 ? 2 : 5), u, setup, ) @@ -1215,11 +1154,10 @@ function interpolate_u_p!(up, u, setup) @kernel function int!(up, u, ::Val{α}, I0) where {α} I = @index(Global, Cartesian) I = I + I0 - up[α][I] = (u[α][I-e(α)] + u[α][I]) / 2 + up[I, α] = (u[I-e(α), α] + u[I, α]) / 2 end for α = 1:D - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) int!(backend, workgroupsize)(up, u, Val(α), I0; ndrange = Np) end up @@ -1246,8 +1184,7 @@ function interpolate_ω_p!(::Dimension{2}, ωp, ω, setup) I = I + I0 ωp[I] = (ω[I-e(1)-e(2)] + ω[I]) / 2 end - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) int!(backend, workgroupsize)(ωp, ω, I0; ndrange = Np) ωp end @@ -1263,10 +1200,9 @@ function interpolate_ω_p!(::Dimension{3}, ωp, ω, setup) I = I + I0 α₊ = mod1(α + 1, D) α₋ = mod1(α - 1, D) - ωp[α][I] = (ω[α][I-e(α₊)-e(α₋)] + ω[α][I]) / 2 + ωp[I, α] = (ω[I-e(α₊)-e(α₋), α] + ω[I, α]) / 2 end - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) for α = 1:D int!(backend, workgroupsize)(ωp, ω, Val(α), I0; ndrange = Np) end @@ -1300,27 +1236,26 @@ function Dfield!(d, G, p, setup; ϵ = eps(eltype(p))) I = I + I0 g = zero(eltype(p)) for α = 1:D - g += (G[α][I-e(α)] + G[α][I])^2 + g += (G[I-e(α), α] + G[I, α])^2 end lap = zero(eltype(p)) # for α = 1:D - # lap += (G[α][I] - G[α][I-e(α)]) / Δ[α][I[α]] + # lap += (G[I, α] - G[I-e(α), α]) / Δ[α][I[α]] # end if D == 2 - lap += (G[1][I] - G[1][I-e(1)]) / Δ[1][I[1]] - lap += (G[2][I] - G[2][I-e(2)]) / Δ[2][I[2]] + lap += (G[I, 1] - G[I-e(1), 1]) / Δ[1][I[1]] + lap += (G[I, 2] - G[I-e(2), 2]) / Δ[2][I[2]] elseif D == 3 - lap += (G[1][I] - G[1][I-e(1)]) / Δ[1][I[1]] - lap += (G[2][I] - G[2][I-e(2)]) / Δ[2][I[2]] - lap += (G[3][I] - G[3][I-e(3)]) / Δ[3][I[3]] + lap += (G[I, 1] - G[I-e(1), 1]) / Δ[1][I[1]] + lap += (G[I, 2] - G[I-e(2), 2]) / Δ[2][I[2]] + lap += (G[I, 3] - G[I-e(3), 3]) / Δ[3][I[3]] end lap = lap > 0 ? max(lap, ϵ) : min(lap, -ϵ) # lap = abs(lap) d[I] = sqrt(g) / 2 / lap end pressuregradient!(G, p, setup) - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) D!(backend, workgroupsize)(d, G, p, I0; ndrange = Np) d end @@ -1352,13 +1287,12 @@ function Qfield!(Q, u, setup) q = zero(eltype(Q)) for α = 1:D, β = 1:D q -= - (u[α][I] - u[α][I-e(β)]) / Δ[β][I[β]] * (u[β][I] - u[β][I-e(α)]) / + (u[I, α] - u[I-e(β), α]) / Δ[β][I[β]] * (u[I, β] - u[I-e(α), β]) / Δ[α][I[α]] / 2 end Q[I] = q end - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) Q!(backend, workgroupsize)(Q, u, I0; ndrange = Np) Q end @@ -1427,7 +1361,7 @@ function kinetic_energy!(ke, u, setup; interpolate_first = false) I = I + I0 k = zero(eltype(ke)) for α = 1:D - k += (u[α][I] + u[α][I-e(α)])^2 + k += (u[I, α] + u[I-e(α), α])^2 end k = k / 8 ke[I] = k @@ -1437,14 +1371,13 @@ function kinetic_energy!(ke, u, setup; interpolate_first = false) I = I + I0 k = zero(eltype(ke)) for α = 1:D - k += u[α][I]^2 + u[α][I-e(α)]^2 + k += u[I, α]^2 + u[I-e(α), α]^2 end k = k / 4 ke[I] = k end ke! = interpolate_first ? efirst! : elast! - I0 = first(Ip) - I0 -= oneunit(I0) + I0 = getoffset(Ip) ke!(backend, workgroupsize)(ke, u, I0; ndrange = Np) ke end @@ -1475,7 +1408,7 @@ function get_scale_numbers(u, setup) (; grid, Re) = setup (; dimension, Iu, Ip, Δ, Δu) = grid D = dimension() - T = eltype(u[1]) + T = eltype(u) visc = 1 / Re Ω = scalewithvolume!(fill!(scalarfield(setup), 1), setup) uavg = @@ -1485,8 +1418,9 @@ function get_scale_numbers(u, setup) D, ) Ωu = .*(Δα...) - field = @. u[α]^2 * Ωu - sum(field[Iu[α]]) / sum(Ωu[Iu[α]]) + uα = eachslice(u; dims = ndims(u)) + field = @. u^2 * Ωu + sum(field[Iu[1], :]) / sum(Ωu[Iu[1]]) end |> sqrt ϵ = dissipation_from_strain(u, setup) ϵ = sum((Ω.*ϵ)[Ip]) / sum(Ω[Ip]) diff --git a/src/pressure.jl b/src/pressure.jl index cdc74f1fc..b96c8d1d1 100644 --- a/src/pressure.jl +++ b/src/pressure.jl @@ -28,9 +28,6 @@ field, resulting in same order pressure as velocity. Differentiable version. """ function pressure(u, temp, t, setup; psolver) - (; grid) = setup - (; dimension, Iu, Ip) = grid - D = dimension() F = momentum(u, temp, t, setup) F = apply_bc_u(F, t, setup; dudt = true) div = divergence(F, setup) @@ -42,9 +39,6 @@ end "Compute pressure from velocity field (in-place version)." function pressure!(p, u, temp, t, setup; psolver, F) - (; grid) = setup - (; dimension, Iu, Ip) = grid - D = dimension() momentum!(F, u, temp, t, setup) apply_bc_u!(F, t, setup; dudt = true) divergence!(p, F, setup) @@ -56,7 +50,7 @@ end "Project velocity field onto divergence-free space (differentiable version)." function project(u, setup; psolver) - T = eltype(u[1]) + T = eltype(u) # Divergence of tentative velocity field div = divergence(u, setup) @@ -73,7 +67,7 @@ end "Project velocity field onto divergence-free space (in-place version)." function project!(u, setup; psolver, p) - T = eltype(u[1]) + T = eltype(u) # Divergence of tentative velocity field divergence!(p, u, setup) @@ -193,9 +187,7 @@ end # Preconditioner function create_laplace_diag(setup) (; grid, workgroupsize) = setup - (; dimension, Δ, Δu, N, Np, Ip) = grid - D = dimension() - δ = Offset{D}() + (; Δ, Δu, Np, Ip) = grid @kernel function _laplace_diag!(z, p, I0) I = @index(Global, Cartesian) I = I + I0 diff --git a/src/processors.jl b/src/processors.jl index e83219831..421977198 100644 --- a/src/processors.jl +++ b/src/processors.jl @@ -64,7 +64,7 @@ timelogger(; showiter && push!(msg, "Iteration $n") showt && push!(msg, @sprintf("t = %g", t)) showdt && push!(msg, @sprintf("Δt = %.2g", Δt)) - showmax && push!(msg, @sprintf("umax = %.2g", maximum(maximum.(abs, u)))) + showmax && push!(msg, @sprintf("umax = %.2g", maximum(abs, u))) showspeed && push!(msg, @sprintf("itertime = %.2g", itertime)) @info join(msg, "\t") end @@ -84,12 +84,11 @@ function observefield( (; dimension, Ip) = setup.grid (; u, temp, t) = state[] D = dimension() - T = eltype(u[1]) # Initialize buffers _f = if fieldname in (1, 2, 3) up = interpolate_u_p(u, setup) - up[fieldname] + upf = selectdim(up, ndims(up), fieldname) elseif fieldname == :velocity up = interpolate_u_p(u, setup) elseif fieldname == :velocitynorm @@ -131,22 +130,30 @@ function observefield( else error("Unknown fieldname") end - _f = _f isa Tuple ? map(f -> Array(f)[Ip], _f) : Array(_f)[Ip] + if ndims(_f) == D + 1 + _f = Array(_f)[Ip, :] + elseif ndims(_f) == D + _f = Array(_f)[Ip] + else + error() + end # Observe field field = lift(state) do (; u, temp, t) f = if fieldname in (1, 2, 3) interpolate_u_p!(up, u, setup) - up[fieldname] + upf elseif fieldname == :velocity interpolate_u_p!(up, u, setup) elseif fieldname == :velocitynorm interpolate_u_p!(up, u, setup) # map((u, v, w) -> √sum(u^2 + v^2 + w^2), up...) if D == 2 - @. upnorm = sqrt(up[1]^2 + up[2]^2) + uptuple = eachslice(up; dims = ndims(up)) + @. upnorm = sqrt(uptuple[1]^2 + uptuple[2]^2) elseif D == 3 - @. upnorm = sqrt(up[1]^2 + up[2]^2 + up[3]^2) + uptuple = eachslice(up; dims = ndims(up)) + @. upnorm = sqrt(uptuple[1]^2 + uptuple[2]^2 + uptuple[3]^2) end elseif fieldname == :vorticity apply_bc_u!(u, t, setup) @@ -179,12 +186,12 @@ function observefield( elseif fieldname == :temperature temp end - if _f isa Tuple - for (_f, f) in zip(_f, f) - copyto!(_f, view(f, Ip)) - end - else + if ndims(f) == D + 1 + copyto!(_f, view(f, Ip, :)) + elseif ndims(f) == D copyto!(_f, view(f, Ip)) + else + error() end _f end @@ -198,7 +205,6 @@ function snapshotsaver(state; setup, fieldnames = (:velocity,), psolver = nothin state isa Observable || (state = Observable(state)) (; grid) = setup (; dimension, xp, Ip) = grid - D = dimension() xparr = getindex.(Array.(xp), Ip.indices) fields = map(fieldname -> observefield(state; setup, fieldname, psolver), fieldnames) @@ -273,7 +279,7 @@ fieldsaver(; setup, nupdate = 1) = on(state) do state state.n % nupdate == 0 || return state = adapt(Array, state) - state.u[1] isa Array && (state = deepcopy(state)) + state.u isa Array && (state = deepcopy(state)) push!(states, state) end states @@ -546,7 +552,7 @@ Create energy history plot. """ function energy_history_plot(state; setup) @assert state isa Observable "Energy history requires observable state." - (; Δ, Ip) = setup.grid + (; Ip) = setup.grid e = scalarfield(setup) _points = Point2f[] points = lift(state) do (; u, t) @@ -578,7 +584,7 @@ function observespectrum(state; setup, npoint = 100, a = typeof(setup.Re)(1 + sq # interpolate_u_p!(up, u, setup) up = u # TODO: Maybe preallocate e and A * e - e = sum(up) do u + e = sum(eachslice(up; dims = D + 1)) do u copyto!(uhat, view(u, Ip)) fft!(uhat) uhathalf = view(uhat, ntuple(α -> 1:K[α], D)...) diff --git a/src/solver.jl b/src/solver.jl index 88c2e18a1..a82a85670 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -32,7 +32,7 @@ function solve_unsteady(; # Cache arrays for intermediate computations cache = ode_method_cache(method, setup), ) - docopy && (ustart = copy.(ustart)) + docopy && (ustart = copy(ustart)) docopy && !isnothing(tempstart) && (tempstart = copy(tempstart)) tstart, tend = tlims @@ -107,14 +107,14 @@ function get_cfl_timestep!(buf, u, setup) Δt = eltype(u[1])(Inf) # Check maximum step size in each dimension - for α = 1:D + for (α, uα) in enumerate(eachslice(u; dims = D + 1)) # Diffusion Δαmin = minimum(view(Δu[α], Iu[α].indices[α])) Δt_diff = Re * Δαmin^2 / 2 # Convection Δα = reshape(Δu[α], ntuple(Returns(1), α - 1)..., :) - @. buf = Δα / abs(u[α]) + @. buf = Δα / abs(uα) Δt_conv = minimum(view(buf, Iu[α])) # Update time step diff --git a/src/time_steppers/isexplicit.jl b/src/time_steppers/isexplicit.jl deleted file mode 100644 index 7a724a430..000000000 --- a/src/time_steppers/isexplicit.jl +++ /dev/null @@ -1,12 +0,0 @@ -""" - isexplicit(method) - -Return `true` if `method` is explicit, i.e. the value at a certain time step is -given explicitly as a function of the previous time steps only. -""" -function isexplicit end - -# By default, not explicit -isexplicit(::AbstractODEMethod) = false -isexplicit(::ExplicitRungeKuttaMethod) = true -isexplicit(::OneLegMethod) = true diff --git a/src/time_steppers/lambda_max.jl b/src/time_steppers/lambda_max.jl deleted file mode 100644 index 8e81500e5..000000000 --- a/src/time_steppers/lambda_max.jl +++ /dev/null @@ -1,20 +0,0 @@ -""" - lambda_diff_max(method) - -Get maximum value of stability region for the diffusion operator. -""" -function lambda_diff_max end - -lambda_diff_max(::AbstractODEMethod) = 1 -lambda_diff_max((; β)::OneLegMethod) = 4 * β / (2 * β + 1) - -""" - lambda_conv_max(method) - - -Get maximum value of stability region for the convection operator (not a very good -indication for the methods that do not include the imaginary axis) -""" -function lambda_conv_max end - -lambda_conv_max(::AbstractODEMethod) = 1 diff --git a/src/time_steppers/step_explicit_runge_kutta.jl b/src/time_steppers/step_explicit_runge_kutta.jl index 641cb207a..89e6193e9 100644 --- a/src/time_steppers/step_explicit_runge_kutta.jl +++ b/src/time_steppers/step_explicit_runge_kutta.jl @@ -3,18 +3,16 @@ create_stepper(::ExplicitRungeKuttaMethod; setup, psolver, u, temp, t, n = 0) = function timestep!(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing, cache) (; setup, psolver, u, temp, t, n) = stepper - (; grid, closure_model, temperature) = setup - (; dimension, Iu) = grid + (; closure_model, temperature) = setup (; A, b, c) = method - (; u₀, ku, p, temp₀, ktemp, diff) = cache - D = dimension() + (; ustart, ku, p, tempstart, ktemp, diff) = cache nstage = length(b) m = closure_model # Update current solution - t₀ = t - copyto!.(u₀, u) - isnothing(temp) || copyto!(temp₀, temp) + tstart = t + copyto!(ustart, u) + isnothing(temp) || copyto!(tempstart, temp) for i = 1:nstage # Compute force at current stage i @@ -28,21 +26,18 @@ function timestep!(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing, end # Add closure term - isnothing(m) || map((k, m) -> k .+= m, ku[i], m(u, θ)) + isnothing(m) || (ku[i] .+= m(u, θ)) # Intermediate time step - t = t₀ + c[i] * Δt + t = tstart + c[i] * Δt # Apply stage forces - for α = 1:D - u[α] .= u₀[α] - for j = 1:i - @. u[α] += Δt * A[i, j] * ku[j][α] - # @. u[α][Iu[α]] += Δt * A[i, j] * ku[j][α][Iu[α]] - end + u .= ustart + for j = 1:i + @. u += Δt * A[i, j] * ku[j] end if !isnothing(temp) - temp .= temp₀ + temp .= tempstart for j = 1:i @. temp += Δt * A[i, j] * ktemp[j] end @@ -65,16 +60,14 @@ end function timestep(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing) (; setup, psolver, u, temp, t, n) = stepper - (; grid, closure_model, temperature) = setup - (; dimension) = grid + (; closure_model, temperature) = setup (; A, b, c) = method - D = dimension() nstage = length(b) m = closure_model # Update current solution (does not depend on previous step size) - t₀ = t - u₀ = u + tstart = t + ustart = u ku = () ktemp = () @@ -89,23 +82,22 @@ function timestep(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing) end # Add closure term - isnothing(m) || (F = F .+ m(u, θ)) + isnothing(m) || (F = F + m(u, θ)) # Store right-hand side of stage i ku = (ku..., F) isnothing(temp) || (ktemp = (ktemp..., Ftemp)) # Intermediate time step - t = t₀ + c[i] * Δt + t = tstart + c[i] * Δt # Apply stage forces - u = u₀ + u = ustart for j = 1:i u = @. u + Δt * A[i, j] * ku[j] - # u = tupleadd(u, @.(Δt * A[i, j] * ku[j])) end if !isnothing(temp) - temp = temp₀ + temp = tempstart for j = 1:i temp = @. temp + Δt * A[i, j] * ktemp[j] end diff --git a/src/time_steppers/step_lmwray3.jl b/src/time_steppers/step_lmwray3.jl index f3f505af3..5ced422a7 100644 --- a/src/time_steppers/step_lmwray3.jl +++ b/src/time_steppers/step_lmwray3.jl @@ -3,12 +3,10 @@ create_stepper(::LMWray3; setup, psolver, u, temp, t, n = 0) = function timestep!(method::LMWray3, stepper, Δt; θ = nothing, cache) (; setup, psolver, u, temp, t, n) = stepper - (; grid, closure_model, temperature) = setup - (; dimension, Iu) = grid + (; closure_model, temperature) = setup (; ustart, ku, p, tempstart, ktemp, diff) = cache - D = dimension() m = closure_model - T = eltype(u[1]) + T = eltype(u) # We wrap the state in x = (; u, temp), and define some # functions that operate on x @@ -26,7 +24,7 @@ function timestep!(method::LMWray3, stepper, Δt; θ = nothing, cache) end # Add closure term - isnothing(m) || map((du, m) -> du .+= m, dx.u, m(u, θ)) + isnothing(m) || (dx.u .+= m(u, θ)) dx end @@ -42,16 +40,14 @@ function timestep!(method::LMWray3, stepper, Δt; θ = nothing, cache) # Copy state x to y function state_copyto!(y, x) - copyto!.(y.u, x.u) + copyto!(y.u, x.u) isnothing(temp) || copyto!(y.temp, x.temp) y end # Compute y = a * x + y for states x, y function state_axpy!(a, x, y) - for α = 1:D - @. y.u[α] += a * x.u[α] - end + @. y.u += a * x.u if !isnothing(temp) @. y.temp += a * x.temp end diff --git a/src/time_steppers/time_stepper_caches.jl b/src/time_steppers/time_stepper_caches.jl index a2b1abc8a..1f0a4b4fa 100644 --- a/src/time_steppers/time_stepper_caches.jl +++ b/src/time_steppers/time_stepper_caches.jl @@ -32,20 +32,20 @@ function ode_method_cache(::OneLegMethod, setup) end function ode_method_cache(method::ExplicitRungeKuttaMethod, setup) - u₀ = vectorfield(setup) + ustart = vectorfield(setup) ns = length(method.b) ku = map(i -> vectorfield(setup), 1:ns) p = scalarfield(setup) if isnothing(setup.temperature) - temp₀ = nothing + tempstart = nothing ktemp = nothing diff = nothing else - temp₀ = scalarfield(setup) + tempstart = scalarfield(setup) ktemp = map(i -> scalarfield(setup), 1:ns) diff = vectorfield(setup) end - (; u₀, ku, p, temp₀, ktemp, diff) + (; ustart, ku, p, tempstart, ktemp, diff) end function ode_method_cache(method::ImplicitRungeKuttaMethod{T}, setup, V, p) where {T} @@ -124,7 +124,7 @@ function ode_method_cache(method::ImplicitRungeKuttaMethod{T}, setup, V, p) wher ) end -function ode_method_cache(method::LMWray3, setup) +function ode_method_cache(::LMWray3, setup) ustart = vectorfield(setup) ku = vectorfield(setup) p = scalarfield(setup) diff --git a/test/chainrules.jl b/test/chainrules.jl index fe8e414c3..d5c13938d 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -18,14 +18,14 @@ end lims = T(0), T(1) x = range(lims..., n + 1), range(lims..., n + 1) for bc in (PeriodicBC(), DirichletBC(), SymmetricBC(), PressureBC()) - boundary_conditions = ((bc, bc), (bc, bc)) + boundary_conditions = (bc, bc), (bc, bc) setup = Setup(; x, Re, boundary_conditions, temperature = temperature_equation(; Pr, Ra, Ge, boundary_conditions), ) - u = randn(T, setup.grid.N), randn(T, setup.grid.N) + u = randn(T, setup.grid.N..., 2) p = randn(T, setup.grid.N) temp = randn(T, setup.grid.N) test_rrule_named(apply_bc_u, u, T(0) ⊢ NoTangent(), setup ⊢ NoTangent()) @@ -64,7 +64,7 @@ end ) setup = Setup(; x, boundary_conditions, Re, temperature) psolver = default_psolver(setup) - u = ntuple(α -> randn(T, setup.grid.N), D) + u = randn(T, setup.grid.N..., D) p = randn(T, setup.grid.N) temp = randn(T, setup.grid.N) (; setup, psolver, u, p, temp) diff --git a/test/operators.jl b/test/operators.jl index 3b5e011b1..5d6217680 100644 --- a/test/operators.jl +++ b/test/operators.jl @@ -34,9 +34,9 @@ end @testitem "Pressure gradient" setup = [Setup2D, Setup3D] begin using Random for setup in (Setup2D.setup, Setup3D.setup) - (; Iu, Ip, dimension) = setup.grid + (; Iu, Ip, Δu, Δ, dimension) = setup.grid D = dimension() - v = randn!.(vectorfield(setup)) + v = randn!(vectorfield(setup)) p = randn!(scalarfield(setup)) T = eltype(p) v = apply_bc_u(v, T(0), setup) @@ -46,23 +46,19 @@ end ΩDv = scalewithvolume(Dv, setup) pDv = sum((p.*ΩDv)[Ip]) vGp = if D == 2 - vGpx = v[1] .* setup.grid.Δu[1] .* setup.grid.Δ[2]' .* Gp[1] - vGpy = v[2] .* setup.grid.Δ[1] .* setup.grid.Δu[2]' .* Gp[2] + vGpx = v[:, :, 1] .* Δu[1] .* Δ[2]' .* Gp[:, :, 1] + vGpy = v[:, :, 2] .* Δ[1] .* Δu[2]' .* Gp[:, :, 2] sum(vGpx[Iu[1]]) + sum(vGpy[Iu[2]]) elseif D == 3 vGpx = - v[1] .* setup.grid.Δu[1] .* reshape(setup.grid.Δ[2], 1, :) .* - reshape(setup.grid.Δ[3], 1, 1, :) .* Gp[1] + v[:, :, :, 1] .* Δu[1] .* reshape(Δ[2], 1, :) .* reshape(Δ[3], 1, 1, :) .* Gp[:, :, :, 1] vGpy = - v[2] .* setup.grid.Δ[1] .* reshape(setup.grid.Δu[2], 1, :) .* - reshape(setup.grid.Δ[3], 1, 1, :) .* Gp[2] + v[:, :, :, 2] .* Δ[1] .* reshape(Δu[2], 1, :) .* reshape(Δ[3], 1, 1, :) .* Gp[:, :, :, 2] vGpz = - v[3] .* setup.grid.Δ[1] .* reshape(setup.grid.Δ[2], 1, :) .* - reshape(setup.grid.Δu[3], 1, 1, :) .* Gp[3] + v[:, :, :, 3] .* Δ[1] .* reshape(Δ[2], 1, :) .* reshape(Δu[3], 1, 1, :) .* Gp[:, :, :, 3] sum(vGpx[Iu[1]]) + sum(vGpy[Iu[2]]) + sum(vGpz[Iu[3]]) end - @test Gp isa Tuple - @test Gp[1] isa Array{T} + @test Gp isa Array{T} @test pDv ≈ -vGp # Check that D = -G' end end @@ -87,23 +83,22 @@ end @testitem "Convection" setup = [Setup2D, Setup3D] begin for (u, setup) in ((Setup2D.u, Setup2D.setup), (Setup3D.u, Setup3D.setup)) (; Iu, Δ, Δu) = setup.grid - T = eltype(u[1]) + T = eltype(u) c = convection(u, setup) - D = length(u) + D = length(Δ) uCu = if D == 2 - uCux = u[1] .* Δu[1] .* Δ[2]' .* c[1] - uCuy = u[2] .* Δ[1] .* Δu[2]' .* c[2] + uCux = u[:, :, 1] .* Δu[1] .* Δ[2]' .* c[:, :, 1] + uCuy = u[:, :, 2] .* Δ[1] .* Δu[2]' .* c[:, :, 2] sum(uCux[Iu[1]]) + sum(uCuy[Iu[2]]) elseif D == 3 Δu1, Δu2, Δu3 = Δu[1], reshape(Δu[2], 1, :), reshape(Δu[3], 1, 1, :) Δp1, Δp2, Δp3 = Δ[1], reshape(Δ[2], 1, :), reshape(Δ[3], 1, 1, :) - uCux = @. u[1] * Δu1 * Δp2 * Δp3 .* c[1] - uCuy = @. u[2] * Δp1 * Δu2 * Δp3 .* c[2] - uCuz = @. u[3] * Δp1 * Δp2 * Δu3 .* c[3] + uCux = @. u[:, :, :, 1] * Δu1 * Δp2 * Δp3 .* c[:, :, :, 1] + uCuy = @. u[:, :, :, 2] * Δp1 * Δu2 * Δp3 .* c[:, :, :, 2] + uCuz = @. u[:, :, :, 3] * Δp1 * Δp2 * Δu3 .* c[:, :, :, 3] sum(uCux[Iu[1]]) + sum(uCuy[Iu[2]]) + sum(uCuz[Iu[3]]) end - @test c isa Tuple - @test c[1] isa Array{T} + @test c isa Array{T} @test uCu ≈ 0 atol = 1e-12 # Check skew-symmetry end end @@ -111,60 +106,58 @@ end @testitem "Diffusion" setup = [Setup2D, Setup3D] begin for (u, setup) in ((Setup2D.u, Setup2D.setup), (Setup3D.u, Setup3D.setup)) T = eltype(u[1]) - (; Iu, Δ, Δu) = setup.grid + (; dimension, Iu, Δ, Δu) = setup.grid d = diffusion(u, setup) - D = length(u) + D = dimension() uDu = if D == 2 - uDux = u[1] .* Δu[1] .* Δ[2]' .* d[1] - uDuy = u[2] .* Δ[1] .* Δu[2]' .* d[2] + uDux = u[:, :, 1] .* Δu[1] .* Δ[2]' .* d[:, :, 1] + uDuy = u[:, :, 2] .* Δ[1] .* Δu[2]' .* d[:, :, 2] sum(uDux[Iu[1]]) + sum(uDuy[Iu[2]]) elseif D == 3 Δu1, Δu2, Δu3 = Δu[1], reshape(Δu[2], 1, :), reshape(Δu[3], 1, 1, :) Δp1, Δp2, Δp3 = Δ[1], reshape(Δ[2], 1, :), reshape(Δ[3], 1, 1, :) - uDux = @. u[1] * Δu1 * Δp2 * Δp3 .* d[1] - uDuy = @. u[2] * Δp1 * Δu2 * Δp3 .* d[2] - uDuz = @. u[3] * Δp1 * Δp2 * Δu3 .* d[3] + uDux = @. u[:, :, :, 1] * Δu1 * Δp2 * Δp3 .* d[:, :, :, 1] + uDuy = @. u[:, :, :, 2] * Δp1 * Δu2 * Δp3 .* d[:, :, :, 2] + uDuz = @. u[:, :, :, 3] * Δp1 * Δp2 * Δu3 .* d[:, :, :, 3] sum(uDux[Iu[1]]) + sum(uDuy[Iu[2]]) + sum(uDuz[Iu[3]]) end - @test d isa Tuple - @test d[1] isa Array{T} + @test d isa Array{T} @test uDu ≤ 0 # Check negativity (dissipation) end end @testitem "Convection-Diffusion" setup = [Setup2D, Setup3D] begin for (u, setup) in ((Setup2D.u, Setup2D.setup), (Setup3D.u, Setup3D.setup)) - cd = IncompressibleNavierStokes.convectiondiffusion!(zero.(u), u, setup) + cd = IncompressibleNavierStokes.convectiondiffusion!(zero(u), u, setup) c = convection(u, setup) d = diffusion(u, setup) - @test all(cd .≈ c .+ d) + @test cd ≈ c + d end end @testitem "Momentum" setup = [Setup2D, Setup3D] begin for (u, setup) in ((Setup2D.u, Setup2D.setup), (Setup3D.u, Setup3D.setup)) - T = eltype(u[1]) + T = eltype(u) m = momentum(u, nothing, T(1), setup) - @test m isa Tuple - @test m[1] isa Array{T} - @test all(all.(!isnan, m)) + @test m isa Array{T} + @test all(!isnan, m) end end @testitem "Other fields" setup = [Setup2D, Setup3D] begin using Random for (u, setup) in ((Setup2D.u, Setup2D.setup), (Setup3D.u, Setup3D.setup)) - T = eltype(u[1]) - D = length(u) + T = eltype(u) + D = setup.grid.dimension() p = randn!(scalarfield(setup)) ω = vorticity(u, setup) D == 2 && @test ω isa Array{T} - D == 3 && @test ω isa Tuple - @test smagorinsky_closure(setup)(u, 0.1) isa Tuple + D == 3 && @test ω isa Array{T} + @test smagorinsky_closure(setup)(u, 0.1) isa Array{T} @test tensorbasis(u, setup) isa Tuple - @test interpolate_u_p(u, setup) isa Tuple + @test interpolate_u_p(u, setup) isa Array{T} D == 2 && @test interpolate_ω_p(ω, setup) isa Array{T} - D == 3 && @test interpolate_ω_p(ω, setup) isa Tuple + D == 3 && @test interpolate_ω_p(ω, setup) isa Array{T} @test Dfield(p, setup) isa Array{T} @test Qfield(u, setup) isa Array{T} D == 2 && @test_throws AssertionError eig2field(u, setup) diff --git a/test/postprocess.jl b/test/postprocess.jl index 80f8f0254..eba8a9148 100644 --- a/test/postprocess.jl +++ b/test/postprocess.jl @@ -43,13 +43,13 @@ @testset "Field saver" begin @test length(outputs.field) == nprocess - @test outputs.field[1].u isa Tuple + @test outputs.field[1].u isa Array @test outputs.field[1].t isa Float64 # Test that different copies are stored i = 1 ii = setup.grid.Iu[i] - a = outputs.field[1].u[i][ii] - b = outputs.field[end].u[i][ii] + a = outputs.field[1].u[ii, i] + b = outputs.field[end].u[ii, i] @test norm(a - b) / norm(b) > 0.05 end