test: add unit tests for newton_div

Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>

Author: Mikolaj-A-Kowalski

test/runtests.jl

@@ -52,6 +52,7 @@ CUDA.allowscalar() do
             include("test_stokes_drift.jl")
             include("test_utils.jl")
             include("test_schedules.jl")
+            include("test_newton_div.jl")
         end
     end
 

test/test_newton_div.jl

@@ -0,0 +1,74 @@
+include("dependencies_for_runtests.jl")
+
+using CUDA
+
+# Generate some random points in a single binade [1;2) interval
+function test_data_in_single_binade(::Type{FT}, size) where {FT}
+    prng = Random.Xoshiro(44)
+    return rand(prng, FT, size) .+ 1.0
+end
+
+@testset "CPU newton_div: $FT $WCT" for (FT, WCT) in Iterators.product((Float32, Float64),
+                                                            (Oceananigans.Utils.NormalDivision,
+                                                             Oceananigans.Utils.ConvertingDivision{Float32}))
+    test_input = test_data_in_single_binade(FT, 1024)
+
+    ref = similar(test_input)
+    output = similar(test_input)
+
+    ref .= FT(π) ./ test_input
+    output .= Oceananigans.Utils.newton_div.(WCT, FT(π), test_input)
+
+    @test isapprox(ref, output)
+end
+
+if CUDA.functional()
+
+    @testset "CUDA newton_div: $FT" for FT in (Float32, Float64)
+        test_input = CuArray(test_data_in_single_binade(FT, 1024))
+
+        WCT = Oceananigans.Utils.BackendOptimizedDivision
+
+        ref = similar(test_input)
+        output = similar(test_input)
+
+        ref .= FT(π) ./ test_input
+        output .= Oceananigans.Utils.newton_div.(WCT, FT(π), test_input)
+
+        @test isapprox(ref, output)
+    end
+end
+
+
+function append_weight_computation_type!(list, weno::WENO{<:Any, <:Any, WCT}) where {WCT}
+    push!(list, WCT)
+    append_weight_computation_type!(list, weno.buffer_scheme)
+end
+append_weight_computation_type!(::Any, ::Any) = nothing
+
+# Extract all weight computation types from WENO
+# Assumes a non-weno buffer scheme will not have WENO buffer scheme
+function get_weight_computation_from_weno_advection(weno::WENO)
+    weight_computation_types = DataType[]
+    append_weight_computation_type!(weight_computation_types, weno)
+    return weight_computation_types
+end
+
+@testset "Verify WENO schemes construction" begin
+
+    # WENO
+    weno5 = WENO(order=7; weight_computation=Oceananigans.Utils.NormalDivision)
+    weight_computation_types = get_weight_computation_from_weno_advection(weno5)
+    @test all(weight_computation_types .== Oceananigans.Utils.NormalDivision)
+
+    # Vector Invariant WENO
+    vector_weno = WENOVectorInvariant(order=9, weight_computation=Oceananigans.Utils.BackendOptimizedDivision)
+
+    for field_name in fieldnames(typeof(vector_weno))
+        field = getfield(vector_weno, field_name)
+        if field isa WENO
+            weight_computation_types = get_weight_computation_from_weno_advection(field)
+            @test all(weight_computation_types .== Oceananigans.Utils.BackendOptimizedDivision)
+        end
+    end
+end