Add topography types

Author: nefrathenrici

docs/src/api.md

@@ -59,6 +59,17 @@ ClimaAtmos.AutoDenseJacobian
 ClimaAtmos.AutoSparseJacobian
 ```
 
+## Topography
+
+```@docs
+ClimaAtmos.CosineTopography
+ClimaAtmos.AgnesiTopography
+ClimaAtmos.ScharTopography
+ClimaAtmos.EarthTopography
+ClimaAtmos.DCMIP200Topography
+ClimaAtmos.Hughes2023Topography
+```
+
 ### Internals
 
 ```@docs

src/solver/type_getters.jl

@@ -429,36 +429,37 @@ function get_initial_condition(parsed_args, atmos)
     end
 end
 
+function get_topography(FT, parsed_args)
+    topo_str = parsed_args["topography"]
+    topo_types = Dict("NoWarp" => NoTopography(),
+        "Cosine2D" => CosineTopography{2, FT}(),
+        "Cosine3D" => CosineTopography{3, FT}(),
+        "Agnesi" => AgnesiTopography{FT}(),
+        "Schar" => ScharTopography{FT}(),
+        "Earth" => EarthTopography(),
+        "Hughes2023" => Hughes2023Topography(),
+        "DCMIP200" => DCMIP200Topography(),
+    )
+
+    @assert topo_str in keys(topo_types)
+    return topo_types[topo_str]
+end
+
 function get_steady_state_velocity(params, Y, parsed_args)
     parsed_args["check_steady_state"] || return nothing
     parsed_args["initial_condition"] == "ConstantBuoyancyFrequencyProfile" &&
         parsed_args["mesh_warp_type"] == "Linear" ||
         error("The steady-state velocity can currently be computed only for a \
                ConstantBuoyancyFrequencyProfile with Linear mesh warping")
-    topography = parsed_args["topography"]
-    steady_state_velocity = if topography == "NoWarp"
-        steady_state_velocity_no_warp
-    elseif topography == "Cosine2D"
-        steady_state_velocity_cosine_2d
-    elseif topography == "Cosine3D"
-        steady_state_velocity_cosine_3d
-    elseif topography == "Agnesi"
-        steady_state_velocity_agnesi
-    elseif topography == "Schar"
-        steady_state_velocity_schar
-    else
-        error("The steady-state velocity for $topography topography cannot \
-               be computed analytically")
-    end
+    topo = get_topography(eltype(params), parsed_args)
     top_level = Spaces.nlevels(axes(Y.c)) + Fields.half
     z_top = Fields.level(Fields.coordinate_field(Y.f).z, top_level)
 
-    # TODO: This can be very expensive! It should be moved to a separate CI job.
     @info "Approximating steady-state velocity"
     s = @timed_str begin
-        ᶜu = steady_state_velocity.(params, Fields.coordinate_field(Y.c), z_top)
+        ᶜu = steady_state_velocity.(topo, params, Fields.coordinate_field(Y.c), z_top)
         ᶠu =
-            steady_state_velocity.(params, Fields.coordinate_field(Y.f), z_top)
+            steady_state_velocity.(topo, params, Fields.coordinate_field(Y.f), z_top)
     end
     @info "Steady-state velocity approximation completed: $s"
     return (; ᶜu, ᶠu)
@@ -813,6 +814,7 @@ function get_simulation(config::AtmosConfig)
     end
 
     tracers = get_tracers(config.parsed_args)
+
     steady_state_velocity =
         get_steady_state_velocity(params, Y, config.parsed_args)
 

src/topography/steady_state_solutions.jl

@@ -1,5 +1,4 @@
 import StaticArrays: @SVector
-
 background_u(::Type{FT}) where {FT} = FT(10)
 background_N(::Type{FT}) where {FT} = FT(0.01) # This needs to be a small value.
 background_T_sfc(::Type{FT}) where {FT} = FT(288)
@@ -167,29 +166,26 @@ end
 ## Steady-state solutions for periodic topography
 ##
 
-function steady_state_velocity_no_warp(params, coord, z_top)
+function steady_state_velocity(::NoTopography, params, coord, z_top)
     FT = eltype(params)
     u = background_u(FT)
     return UVW(u, FT(0), FT(0))
 end
 
-function steady_state_velocity_cosine_2d(params, coord, z_top)
+function steady_state_velocity(t::CosineTopography{2}, params, coord, z_top)
     FT = eltype(params)
-    (; λ) = cosine_params(FT)
     (; x, z) = coord
-    return steady_state_velocity_cosine(params, x, FT(0), z, λ, FT(Inf), z_top)
+    return steady_state_velocity_cosine(params, x, FT(0), z, t.λ, oftype(t.λ, Inf), z_top, t.h_max)
 end
 
-function steady_state_velocity_cosine_3d(params, coord, z_top)
+function steady_state_velocity(t::CosineTopography{3}, params, coord, z_top)
     FT = eltype(params)
-    (; λ) = cosine_params(FT)
     (; x, y, z) = coord
-    return steady_state_velocity_cosine(params, x, y, z, λ, λ, z_top)
+    return steady_state_velocity_cosine(params, x, y, z, t.λ, t.λ, z_top, t.h_max)
 end
 
-function steady_state_velocity_cosine(params, x, y, z, λ_x, λ_y, z_top)
+function steady_state_velocity_cosine(params, x, y, z, λ_x, λ_y, z_top, h_max)
     FT = eltype(params)
-    (; h_max) = cosine_params(FT)
     u = background_u(FT)
     h = topography_cosine(x, y, λ_x, λ_y, h_max)
     η = (z - h) / (1 - h / z_top)
@@ -253,14 +249,14 @@ function steady_state_velocity_mountain_2d(
     return UVW(u + Δu, Δv, Δw)
 end
 
-function steady_state_velocity_agnesi(params, coord, z_top)
+function steady_state_velocity(t::AgnesiTopography, params, coord, z_top)
     FT = eltype(params)
-    (; h_max, x_center, a) = agnesi_params(FT)
+    (; h_max, x_center, a) = t
     topography_agnesi_Fh(k_x) = h_max * a / 2 * exp(-a * abs(k_x))
     n_efolding_intervals = -log(eps(FT))
     k_x_max = n_efolding_intervals / a
     return steady_state_velocity_mountain_2d(
-        topography_agnesi,
+        topography_function(t),
         topography_agnesi_Fh,
         params,
         coord,
@@ -270,9 +266,9 @@ function steady_state_velocity_agnesi(params, coord, z_top)
     )
 end
 
-function steady_state_velocity_schar(params, coord, z_top)
+function steady_state_velocity(t::ScharTopography, params, coord, z_top)
     FT = eltype(params)
-    (; h_max, x_center, λ, a) = schar_params(FT)
+    (; h_max, x_center, λ, a) = t
     k_peak = 2 * FT(π) / λ
     Fh_coef = h_max * a / (8 * sqrt(FT(π)))
     topography_schar_Fh(k_x) =
@@ -284,7 +280,7 @@ function steady_state_velocity_schar(params, coord, z_top)
     n_efolding_intervals = -log(eps(FT))
     k_x_max = k_peak + 2 * sqrt(n_efolding_intervals) / a
     return steady_state_velocity_mountain_2d(
-        topography_schar,
+        topography_function(t),
         topography_schar_Fh,
         params,
         coord,

src/topography/topography.jl

@@ -1,8 +1,101 @@
+using ClimaCore: Geometry, Spaces, Fields
+
+##
+## Topography profiles for 2D and 3D boxes
+##
+
+# The parameters of these profiles should be defined separately so that they
+# can also be used to compute analytic solutions.
+
+abstract type AbstractTopography end
+Base.broadcastable(t::AbstractTopography) = tuple(t)
+topography_function(topo) = Base.Fix1(topography_function, topo)
+
+struct NoTopography <: AbstractTopography end
+
+# Analytical topography types for idealized test cases
+
+"""
+    CosineTopography{D, FT}(; h_max = 25, λ = 25e3)
+    
+Cosine hill topography in 2D or 3D.
+
+# Arguments
+- `h_max::FT`: Maximum elevation (m)
+- `λ::FT`: Wavelength of the cosine hills (m)
+"""
+Base.@kwdef struct CosineTopography{D, FT} <: AbstractTopography
+    h_max::FT = 25.0
+    λ::FT = 25e3
+end
+
+topography_function(t::CosineTopography{2}, coord) =
+    topography_cosine(coord.x, zero(coord.x), t.λ, oftype(t.λ, Inf), t.h_max)
+
+topography_function(t::CosineTopography{3}, coord) =
+    topography_cosine(coord.x, coord.y, t.λ, t.λ, t.h_max)
+
+topography_cosine(x, y, λ_x, λ_y, h_max) =
+    h_max * cospi(2 * x / λ_x) * cospi(2 * y / λ_y)
+
+"""
+    AgnesiTopography{FT}(; h_max = 25, x_center = 50e3, a = 5e3)
+
+Witch of Agnesi mountain topography for 2D simulations.
+
+# Arguments
+- `h_max`: Maximum elevation (m)
+- `x_center`: Center position (m)
+- `a`: Mountain width parameter (m)
+"""
+Base.@kwdef struct AgnesiTopography{FT} <: AbstractTopography
+    h_max::FT = 25.0
+    x_center::FT = 50e3
+    a::FT = 5e3
+end
+
+topography_function((; h_max, x_center, a)::AgnesiTopography, (; x)) =
+    h_max / (1 + ((x - x_center) / a)^2)
+
+"""
+    ScharTopography{FT}(; h_max = 25, x_center = 50e3, λ = 4e3, a = 5e3)
+
+Schar mountain topography for 2D simulations.
+
+# Arguments
+- `h_max`: Maximum elevation (m)
+- `x_center`: Center position (m)
+- `λ`: Wavelength parameter (m)
+- `a`: Mountain width parameter (m)
+"""
+Base.@kwdef struct ScharTopography{FT} <: AbstractTopography
+    h_max::FT = 25.0
+    x_center::FT = 50e3
+    λ::FT = 4e3
+    a::FT = 5e3
+end
+
+topography_function((; h_max, x_center, λ, a)::ScharTopography, (; x)) =
+    h_max * exp(-(x - x_center)^2 / a^2) * cospi((x - x_center) / λ)^2
+
+# Data-based topography types
+
+"""
+    EarthTopography()
+
+Earth topography from ETOPO2022 data files.
+"""
+struct EarthTopography <: AbstractTopography end
+
 """
-    topography_dcmip200(coord)
+    DCMIP200Topography()
 
 Surface elevation for the DCMIP-2-0-0 test problem.
 """
+struct DCMIP200Topography <: AbstractTopography end
+
+topography_function(::DCMIP200Topography, coord) = topography_dcmip200(coord)
+
 function topography_dcmip200(coord)
     FT = Geometry.float_type(coord)
     λ, ϕ = coord.long, coord.lat
@@ -21,10 +114,14 @@ function topography_dcmip200(coord)
 end
 
 """
-    topography_hughes2023(coord)
+    Hughes2023Topography()
 
 Surface elevation for baroclinic wave test from Hughes and Jablonowski (2023).
 """
+struct Hughes2023Topography <: AbstractTopography end
+
+topography_function(::Hughes2023Topography, coord) = topography_hughes2023(coord)
+
 function topography_hughes2023(coord)
     FT = Geometry.float_type(coord)
     λ, ϕ = coord.long, coord.lat
@@ -50,66 +147,3 @@ function topography_hughes2023(coord)
         ),
     )
 end
-
-##
-## Topography profiles for 2D and 3D boxes
-##
-
-# The parameters of these profiles should be defined separately so that they
-# can also be used to compute analytic solutions.
-
-"""
-    topography_agnesi(coord)
-
-Surface elevation for a 2D Witch of Agnesi mountain, centered at `x = 50 km`.
-"""
-function topography_agnesi(coord)
-    FT = Geometry.float_type(coord)
-    (; x) = coord
-    (; h_max, x_center, a) = agnesi_params(FT)
-    return h_max / (1 + ((x - x_center) / a)^2)
-end
-agnesi_params(::Type{FT}) where {FT} =
-    (; h_max = FT(25), x_center = FT(50e3), a = FT(5e3))
-
-"""
-    topography_schar(coord)
-
-Surface elevation for a 2D Schar mountain, centered at `x = 50 km`.
-"""
-function topography_schar(coord)
-    FT = Geometry.float_type(coord)
-    (; x) = coord
-    (; h_max, x_center, λ, a) = schar_params(FT)
-    return h_max * exp(-(x - x_center)^2 / a^2) * cospi((x - x_center) / λ)^2
-end
-schar_params(::Type{FT}) where {FT} =
-    (; h_max = FT(25), x_center = FT(50e3), λ = FT(4e3), a = FT(5e3))
-
-"""
-    topography_cosine_2d(coord)
-
-Surface elevation for 2D cosine hills.
-"""
-function topography_cosine_2d(coord)
-    FT = Geometry.float_type(coord)
-    (; x) = coord
-    (; h_max, λ) = cosine_params(FT)
-    return topography_cosine(x, FT(0), λ, FT(Inf), h_max)
-end
-
-"""
-    topography_cosine_3d(coord)
-
-Surface elevation for 3D cosine hills.
-"""
-function topography_cosine_3d(coord)
-    FT = Geometry.float_type(coord)
-    (; x, y) = coord
-    (; h_max, λ) = cosine_params(FT)
-    return topography_cosine(x, y, λ, λ, h_max)
-end
-
-topography_cosine(x, y, λ_x, λ_y, h_max) =
-    h_max * cospi(2 * x / λ_x) * cospi(2 * y / λ_y)
-cosine_params(::Type{FT}) where {FT} = (; h_max = FT(25), λ = FT(25e3))