fix: update for arrays

Author: agdestein

lib/NeuralClosure/src/closure.jl

@@ -1,37 +1,14 @@
 """
 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
-end
+wrappedclosure(m) =
+    function neuralclosure(u, θ)
+        s = size(u)
+        u = reshape(u, s..., 1) # Add sample dim
+        mu = m(u, θ)
+        mu = pad_circular(mu, 1)
+        mu = reshape(mu, s) # Remove sample dim
+    end
 
 """
 Create neural closure model from layers.

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])

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
@@ -226,7 +224,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)[it], Iu[α], α),
                 )
             end
             u

lib/NeuralClosure/src/filter.jl

@@ -24,32 +24,31 @@ struct VolumeAverage <: AbstractFilter end
     Φ(vectorfield(setup_les), u, setup_les, compression)
 
 function (::FaceAverage)(v, u, setup_les, comp)
-    (; grid, workgroupsize) = setup_les
+    (; grid, backend, workgroupsize) = setup_les
     (; Nu, Iu) = grid
     D = length(u)
     @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
+    (; grid, boundary_conditions, backend, workgroupsize) = setup_les
+    (; N) = grid
     D = length(u)
     e = Offset(D)
     @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions)
@@ -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,7 +78,7 @@ 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
+    (; grid, boundary_conditions, backend, workgroupsize) = setup_les
     (; N, Nu, Iu) = grid
     D = length(u)
     @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions)
@@ -93,16 +92,15 @@ function (::VolumeAverage)(v, u, setup_les, comp)
             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

src/IncompressibleNavierStokes.jl

@@ -15,7 +15,6 @@ using DocStringExtensions
 using FFTW
 using IterativeSolvers
 using KernelAbstractions
-using KernelAbstractions.Extras.LoopInfo: @unroll
 using LinearAlgebra
 using Makie
 using NNlib
@@ -73,6 +72,8 @@ 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
@@ -94,7 +95,6 @@ export fieldplot, energy_history_plot, energy_spectrum_plot
 
 # Setup
 export Setup, temperature_equation, scalarfield, vectorfield
-export CPU # TODO: This annoys Documenter, need to include this docstring somehow
 
 # 1D grids
 export stretched_grid, cosine_grid, tanh_grid
@@ -109,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