Fix triply periodic pressure solver and add/cleanup some tests

src/BoundaryConditions/show_boundary_conditions.jl

@@ -4,7 +4,7 @@ import Oceananigans.Grids: short_show
 ##### BoundaryCondition
 #####
 
-print_condition(n::Number) = "$n"
+print_condition(n::Union{Nothing, Number}) = "$n"
 print_condition(A::AbstractArray) = "$(Base.dims2string(size(A))) $(typeof(A))"
 print_condition(bf::Union{ParameterizedBoundaryConditionFunction, BoundaryFunction}) = print_condition(bf.func)
 

src/Diagnostics/Diagnostics.jl

@@ -10,7 +10,6 @@ using Oceananigans,
 
 using Oceananigans: AbstractDiagnostic
 
-include("diagnostics_kernels.jl")
 include("nan_checker.jl")
 include("horizontal_average.jl")
 include("time_series.jl")

src/Fields/set!.jl

@@ -1,6 +1,6 @@
 using CUDA
 using KernelAbstractions: @kernel, @index, CUDADevice
-using Oceananigans.Architectures: device
+using Oceananigans.Architectures: device, GPU
 using Oceananigans.Utils: work_layout
 
 """
@@ -84,7 +84,7 @@ end
 
 # Set the GPU field `u` to the CuArray `v`.
 @hascuda function set!(u::AbstractGPUField, v::CuArray)
-    launch!(CUDADevice(), u.grid, :xyz, _set_gpu!, u.data, v, u.grid)
+    launch!(GPU(), u.grid, :xyz, _set_gpu!, u.data, v, u.grid)
     return nothing
 end
 

src/Oceananigans.jl

@@ -119,9 +119,9 @@ include("Architectures.jl")
 
 using Oceananigans.Architectures: @hascuda
 @hascuda begin
-    println("CUDA-enabled GPU(s) detected:")
+    @debug "CUDA-enabled GPU(s) detected:"
     for (gpu, dev) in enumerate(CUDA.devices())
-        println(dev)
+        @debug "$dev: $(CUDA.name(dev))"
     end
 end
 

src/Solvers/discrete_eigenvalues.jl

@@ -68,7 +68,7 @@ equation with periodic boundary conditions in the z-dimension on `grid`.
 function λk(grid, ::PeriodicBC)
     Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
     ks = reshape(1:Nz, 1, 1, Nz)
-    return @. (2sin((ks - 1) * π / 2Nz) / (Lz / Nz))^2
+    return @. (2sin((ks - 1) * π / Nz) / (Lz / Nz))^2
 end
 
 """

src/Solvers/pressure_solver.jl

@@ -26,14 +26,12 @@ function PressureSolver(arch, grid, pressure_bcs, planner_flag=FFTW.PATIENT)
     bc_symbol = Symbol(x, y, z)
 
     if bc_symbol == :PPP
-        @warn "TriplyPeriodicPressureSolver is still untested."
         return TriplyPeriodicPressureSolver(arch, grid, pressure_bcs, planner_flag)
     elseif bc_symbol == :PPN
         return HorizontallyPeriodicPressureSolver(arch, grid, pressure_bcs, planner_flag)
     elseif bc_symbol == :PNN
         return ChannelPressureSolver(arch, grid, pressure_bcs, planner_flag)
     elseif bc_symbol == :NNN
-        @warn "BoxPressureSolver is still untested."
         return BoxPressureSolver(arch, grid, pressure_bcs, planner_flag)
     else
         throw(ArgumentError("Unsupported pressure boundary conditions: $bc_symbol"))

src/Solvers/triply_periodic_pressure_solver.jl

@@ -67,7 +67,7 @@ function TriplyPeriodicPressureSolver(::GPU, grid, pressure_bcs, no_args...)
     IFFTxyz! = plan_backward_transform(storage, x_bc, [1, 2, 3])
     @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, GPU} done!"
 
-    transforms = (FFTxyz! =  FFTxyz!, IFFTxyz! = IFFTxy!)
+    transforms = (FFTxyz! =  FFTxyz!, IFFTxyz! = IFFTxyz!)
 
     return PressureSolver(TriplyPeriodic(), GPU(), wavenumbers, storage, transforms, nothing)
 end

test/runtests.jl

@@ -9,6 +9,7 @@ using CUDA
 using JLD2
 using FFTW
 using OffsetArrays
+using SeawaterPolynomials
 
 using Oceananigans
 using Oceananigans.Architectures
@@ -36,59 +37,15 @@ using TimesDates: TimeDate
 using Statistics: mean
 using LinearAlgebra: norm
 using NCDatasets: Dataset
-
-using SeawaterPolynomials
+using KernelAbstractions: @kernel, @index, Event
 
 import Oceananigans.Fields: interior
-import Oceananigans.Utils: datatuple
+import Oceananigans.Utils: launch!, datatuple
 
-using Oceananigans.Diagnostics: run_diagnostic, velocity_div!
+using Oceananigans.Diagnostics: run_diagnostic
 using Oceananigans.TimeSteppers: _compute_w_from_continuity!
 using Oceananigans.AbstractOperations: Computation, compute!
 
-#####
-##### On CI servers select the GPU with the most available memory or with the
-##### highest capability if testing needs to be thorough).
-##### Source credit: https://github.com/JuliaGPU/CuArrays.jl/pull/526
-#####
-
-@hascuda begin
-    gpu_candidates = [(dev=dev, cap=CUDA.capability(dev),
-                       mem=CUDA.CuContext(ctx -> CUDA.available_memory(), dev))
-                       for dev in CUDA.devices()]
-
-    thorough = parse(Bool, get(ENV, "CI_THOROUGH", "false"))
-    if thorough
-        sort!(gpu_candidates, by=x->(x.cap, x.mem))
-    else
-        sort!(gpu_candidates, by=x->x.mem)
-    end
-
-    pick = last(gpu_candidates)
-    device!(pick.dev)
-end
-
-#####
-##### Useful utilities
-#####
-
-function get_model_field(field_name, model)
-    if field_name ∈ (:u, :v, :w)
-        return getfield(model.velocities, field_name)
-    else
-        return getfield(model.tracers, field_name)
-    end
-end
-
-datatuple(A) = NamedTuple{propertynames(A)}(Array(data(a)) for a in A)
-
-function get_output_tuple(output, iter, tuplename)
-    file = jldopen(output.filepath, "r")
-    output_tuple = file["timeseries/$tuplename/$iter"]
-    close(file)
-    return output_tuple
-end
-
 #####
 ##### Testing parameters
 #####
@@ -113,6 +70,8 @@ closures = (
 ##### Run tests!
 #####
 
+include("runtests_utils.jl")
+
 with_logger(ModelLogger()) do
     @testset "Oceananigans" begin
         include("test_grids.jl")

test/runtests_utils.jl

@@ -0,0 +1,34 @@
+#####
+##### Useful utilities
+#####
+
+function get_model_field(field_name, model)
+    if field_name ∈ (:u, :v, :w)
+        return getfield(model.velocities, field_name)
+    else
+        return getfield(model.tracers, field_name)
+    end
+end
+
+datatuple(A) = NamedTuple{propertynames(A)}(Array(data(a)) for a in A)
+
+function get_output_tuple(output, iter, tuplename)
+    file = jldopen(output.filepath, "r")
+    output_tuple = file["timeseries/$tuplename/$iter"]
+    close(file)
+    return output_tuple
+end
+
+#####
+##### Useful kernels
+#####
+
+@kernel function ∇²!(grid, f, ∇²f)
+    i, j, k = @index(Global, NTuple)
+    @inbounds ∇²f[i, j, k] = ∇²(i, j, k, grid, f)
+end
+
+@kernel function divergence!(grid, u, v, w, div)
+    i, j, k = @index(Global, NTuple)
+    @inbounds div[i, j, k] = divᶜᶜᶜ(i, j, k, grid, u, v, w)
+end

test/test_fields.jl

@@ -68,6 +68,7 @@ end
         @info "  Testing field setting..."
 
         for arch in archs, FT in float_types
+	    ArrayType = array_type(arch)
             grid = RegularCartesianGrid(FT, size=N, extent=L, topology=(Periodic, Periodic, Bounded))
 
             for fieldtype in fieldtypes, val in vals
@@ -76,8 +77,7 @@ end
 
             for fieldtype in fieldtypes
                 field = fieldtype(arch, grid)
-                A = rand(FT, N...)
-                arch isa GPU && (A = CuArray(A))
+                A = rand(FT, N...) |> ArrayType
                 set!(field, A)
                 @test field.data[2, 4, 6] == A[2, 4, 6]
             end