Merge pull request CliMA#438 from climate-machine/more-coriolis-options

More coriolis options

src/Oceananigans.jl

@@ -12,7 +12,7 @@ export
     CPU, GPU,
 
     # Constants
-    FPlane,
+    FPlane, BetaPlane,
     second, minute, hour, day,
 
     # Grids

src/coriolis.jl

@@ -1,45 +1,101 @@
+using .TurbulenceClosures: ▶xy_cfa, ▶xy_fca
+
 #####
 ##### Functions for non-rotating models
 #####
 
-@inline x_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)
-@inline y_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)
-@inline z_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)
+@inline x_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
+@inline y_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
+@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
 
 #####
 ##### The 'FPlane' approximation. This is equivalent to a model with a constant
 ##### rotation rate around its vertical axis.
 #####
 
 """
-    FPlane{T} <: AbstractRotation
+    FPlane{FT} <: AbstractRotation
 
 A parameter object for constant rotation around a vertical axis.
 """
-struct FPlane{T} <: AbstractRotation
-    f :: T
+struct FPlane{FT} <: AbstractRotation
+    f :: FT
 end
 
 """
-    FPlane([T=Float64;] f)
+    FPlane([FT=Float64;] f=nothing, rotation_rate=nothing, latitude=nothing)
 
 Returns a parameter object for constant rotation at the angular frequency
-`2f`, and therefore with background vorticity `f`, around a vertical axis.
+`f/2`, and therefore with background vorticity `f`, around a vertical axis.
+If `f` is not specified, it is calculated from `rotation_rate` and
+`latitude` according to the relation `f = 2*rotation_rate*sind(latitude).
 
-Also called `FPlane`, after the "f-plane" approximation for the local effect of 
+Also called `FPlane`, after the "f-plane" approximation for the local effect of
 Earth's rotation in a planar coordinate system tangent to the Earth's surface.
 """
-function FPlane(T::DataType=Float64; f) 
-    return FPlane{T}(f)
+function FPlane(FT=Float64; f=nothing, rotation_rate=nothing, latitude=nothing)
+
+    if f == nothing && rotation_rate != nothing && latitude != nothing
+        return FPlane{FT}(2rotation_rate*sind(latitude))
+    elseif f != nothing && rotation_rate == nothing && latitude == nothing
+        return FPlane{FT}(f)
+    else
+        throw(ArgumentError("Either both keywords rotation_rate and
+                             latitude must be specified, *or* only f
+                             must be specified."))
+    end
+end
+
+@inline x_f_cross_U(i, j, k, grid, coriolis::FPlane, U) = - coriolis.f * ▶xy_fca(i, j, k, grid, U.v)
+@inline y_f_cross_U(i, j, k, grid, coriolis::FPlane, U) =   coriolis.f * ▶xy_cfa(i, j, k, grid, U.u)
+@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::FPlane, U) where FT = zero(FT)
+
+#####
+##### The Beta Plane
+#####
+
+"""
+    BetaPlane{T} <: AbstractRotation
+
+A parameter object for meridionally increasing Coriolis
+parameter (`f = f₀ + βy`).
+"""
+struct BetaPlane{T} <: AbstractRotation
+    f₀ :: T
+     β :: T
 end
 
-@inline fv(i, j, k, grid::RegularCartesianGrid{T}, f, v) where T = 
-    T(0.5) * f * (avgy_f2c(grid, v, i-1,  j, k) + avgy_f2c(grid, v, i, j, k))
+"""
+    BetaPlane([T=Float64;] f₀=nothing, β=nothing, 
+                           rotation_rate=nothing, latitude=nothing, radius=nothing)
 
-@inline fu(i, j, k, grid::RegularCartesianGrid{T}, f, u) where T = 
-    T(0.5) * f * (avgx_f2c(grid, u, i,  j-1, k) + avgx_f2c(grid, u, i, j, k))
+A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).
+
+The user may specify both `f₀` and `β`, or the three parameters
+`rotation_rate`, `latitude`, and `radius` that specify the rotation rate and radius 
+of a planet, and the central latitude at which the `β`-plane approximation is to be made.
+"""
+function BetaPlane(T=Float64; f₀=nothing, β=nothing, 
+                              rotation_rate=nothing, latitude=nothing, radius=nothing)
 
-@inline x_f_cross_U(i, j, k, grid, rotation::FPlane, U) = -fv(i, j, k, grid, rotation.f, U.v)
-@inline y_f_cross_U(i, j, k, grid, rotation::FPlane, U) =  fu(i, j, k, grid, rotation.f, U.u)
-@inline z_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::FPlane, U) where T = zero(T)
+    f_and_β = f₀ != nothing && β != nothing
+    planet_parameters = rotation_rate != nothing && latitude != nothing && radius != nothing
+
+    (f_and_β && all(Tuple(p === nothing for p in (rotation_rate, latitude, radius)))) || 
+    (planet_parameters && all(Tuple(p === nothing for p in (f₀, β)))) ||
+        throw(ArgumentError("Either both keywords f₀ and β must be specified, 
+                            *or* all of rotation_rate, latitude, and radius."))
+
+    if planet_parameters
+        f₀ = 2rotation_rate * sind(latitude)
+         β = 2rotation_rate * cosd(latitude) / radius
+     end
+
+    return BetaPlane{T}(f₀, β)
+end
 
+@inline x_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
+    @inbounds - (coriolis.f₀ + coriolis.β * grid.yC[j]) * ▶xy_fca(i, j, k, grid, U.v)
+@inline y_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
+    @inbounds   (coriolis.f₀ + coriolis.β * grid.yF[j]) * ▶xy_cfa(i, j, k, grid, U.u)
+@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::BetaPlane, U) where FT = zero(FT)

test/regression_tests/ocean_large_eddy_simulation_regression_test.jl

@@ -63,7 +63,7 @@ function run_ocean_large_eddy_simulation_regression_test(arch, closure)
     # Model instantiation
     model = Model(
              architecture = arch,
-                     grid = RegularCartesianGrid(size=(16, 16, 16), length=(16, 16, 16)),
+                 grid = RegularCartesianGrid(size=(16, 16, 16), length=(16, 16, 16)),
                  coriolis = FPlane(f=1e-4),
                  buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=2e-4, β=8e-4)),
                   closure = closure,

test/test_coriolis.jl

@@ -1,14 +1,43 @@
-function instantiate_coriolis(T)
+function instantiate_fplane_1(T)
     coriolis = FPlane(T, f=π)
     return coriolis.f == T(π)
 end
 
+function instantiate_fplane_2(T)
+    coriolis = FPlane(T, rotation_rate=2, latitude=30)
+    return coriolis.f == T(2)
+end
+
+function instantiate_betaplane_1(T)
+    coriolis = BetaPlane(T, f₀=π, β=2π)
+    return coriolis.f₀ == T(π)
+end
+
+function instantiate_betaplane_2(T)
+    coriolis = BetaPlane(T, latitude=70, radius=2π, rotation_rate=3π)
+    return coriolis.f₀ == T(6π * sind(70))
+end
+
 @testset "Coriolis" begin
     println("Testing Coriolis...")
 
     @testset "Coriolis" begin
         for T in float_types
-            @test instantiate_coriolis(T)
+            @test instantiate_fplane_1(T)
+            @test instantiate_fplane_2(T)
+            @test instantiate_betaplane_1(T)
+            @test instantiate_betaplane_2(T)
+            # Test that FPlane throws an ArgumentError
+            @test_throws ArgumentError FPlane(T, rotation_rate=7e-5)
+            @test_throws ArgumentError FPlane(T, latitude=40)
+            @test_throws ArgumentError FPlane(T, f=1, rotation_rate=7e-5)
+            @test_throws ArgumentError FPlane(T, f=1, latitude=40)
+            @test_throws ArgumentError FPlane(T, f=1, rotation_rate=7e-5, latitude=40)
+            # Non-exhaustively test that BetaPlane throws an ArgumentError
+            @test_throws ArgumentError BetaPlane(T, f₀=1)
+            @test_throws ArgumentError BetaPlane(T, β=1)
+            @test_throws ArgumentError BetaPlane(T, latitude=70)
+            @test_throws ArgumentError BetaPlane(T, f₀=1e-4, β=1e-11, latitude=70)
         end
     end
 end

test/test_regression.jl

@@ -1,3 +1,4 @@
+
 include("regression_tests/thermal_bubble_regression_test.jl")
 include("regression_tests/rayleigh_benard_regression_test.jl")
 include("regression_tests/ocean_large_eddy_simulation_regression_test.jl")

test/test_time_stepping.jl

@@ -5,6 +5,17 @@ function time_stepping_works(arch, FT, Closure)
     return true  # test that no errors/crashes happen when time stepping.
 end
 
+function time_stepping_works_with_beta_plane(arch, FT)
+    model = Model(
+                    float_type = FT,
+                  architecture = arch,
+                          grid = RegularCartesianGrid(FT, size=(16, 16, 16), length=(1, 2, 3)),
+                      coriolis = BetaPlane(FT, f₀=1e-4, β=1e-11)
+                 )
+    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)),
@@ -151,6 +162,10 @@ Closures = (ConstantIsotropicDiffusivity, ConstantAnisotropicDiffusivity,
         @test time_stepping_works(arch, FT, Closure)
     end
 
+    for arch in archs, FT in float_types
+        @test time_stepping_works_with_beta_plane(arch, FT)
+    end
+
     @testset "2nd-order Adams-Bashforth" begin
         println("  Testing 2nd-order Adams-Bashforth...")
         for arch in archs, FT in float_types