Fix and test RoquetIdealizedNonlinearEquationOfState (#574)

Fix and test `RoquetIdealizedNonlinearEquationOfState`

Author: ali-ramadhan

src/Oceananigans.jl

@@ -33,7 +33,7 @@ export
     FPlane, BetaPlane,
 
     # Buoyancy and equations of state
-    BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState,
+    BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState, RoquetIdealizedNonlinearEquationOfState,
 
     # Surface waves via Craik-Leibovich equations
     SurfaceWaves,

src/buoyancy.jl

@@ -230,7 +230,7 @@ const LinearSeawaterBuoyancy = SeawaterBuoyancy{FT, <:LinearEquationOfState} whe
 @inline haline_contractionᶜᶜᶠ(i, j, k, grid, eos, C) = @inbounds haline_contraction(ℑzᵃᵃᶠ(i, j, k, grid, C.T), ℑzᵃᵃᶠ(i, j, k, grid, C.S), Dᵃᵃᶠ(i, j, k, grid), eos)
 
 @inline buoyancy_perturbation(i, j, k, grid, b::AbstractBuoyancy{<:AbstractNonlinearEquationOfState}, C) =
-    - b.gravitational_acceleration * ρ′(i, j, k, grid, b.equation_of_state, C) / b.ρ₀
+    - b.gravitational_acceleration * ρ′(i, j, k, grid, b.equation_of_state, C) / b.equation_of_state.ρ₀
 
 #####
 ##### Roquet et al 2015 idealized nonlinear equations of state

test/test_time_stepping.jl

@@ -8,6 +8,18 @@ function time_stepping_works(arch, FT, Closure)
     return true  # Test that no errors/crashes happen when time stepping.
 end
 
+function time_stepping_works_with_nonlinear_eos(arch, FT, eos_type)
+    grid = RegularCartesianGrid(FT; size=(16, 16, 16), length=(1, 2, 3))
+
+    eos = RoquetIdealizedNonlinearEquationOfState(eos_type)
+    b = SeawaterBuoyancy(equation_of_state=eos)
+
+    model = Model(architecture=arch, float_type=FT, grid=grid, buoyancy=b)
+    time_step!(model, 1, 1)
+
+    return true  # Test that no errors/crashes happen when time stepping.
+end
+
 function run_first_AB2_time_step_tests(arch, FT)
     add_ones(args...) = 1.0
     model = Model(grid=RegularCartesianGrid(FT; size=(16, 16, 16), length=(1, 2, 3)),
@@ -159,6 +171,16 @@ Closures = (ConstantIsotropicDiffusivity, ConstantAnisotropicDiffusivity,
         @test time_stepping_works(arch, FT, Closure)
     end
 
+    @testset "Idealized nonlinear equation of state" begin
+        for arch in archs, FT in float_types
+            for eos_type in keys(Oceananigans.optimized_roquet_coeffs)
+                println("  Testing that time stepping works with " *
+                        "RoquetIdealizedNonlinearEquationOfState [$arch, $FT, $eos_type]")
+                @test time_stepping_works_with_nonlinear_eos(arch, FT, eos_type)
+            end
+        end
+    end
+
     @testset "2nd-order Adams-Bashforth" begin
         println("  Testing 2nd-order Adams-Bashforth...")
         for arch in archs, FT in float_types