Add topography module and types

Author: nefrathenrici

docs/src/api.md

@@ -59,6 +59,23 @@ ClimaAtmos.AutoDenseJacobian
 ClimaAtmos.AutoSparseJacobian
 ```
 
+## Topography
+
+```@docs
+ClimaAtmos.CosineTopography
+ClimaAtmos.AgnesiTopography
+ClimaAtmos.ScharTopography
+ClimaAtmos.EarthTopography
+ClimaAtmos.DCMIP200Topography
+ClimaAtmos.Hughes2023Topography
+ClimaAtmos.topography_schar
+ClimaAtmos.topography_cosine_3d
+ClimaAtmos.topography_agnesi
+ClimaAtmos.topography_hughes2023
+ClimaAtmos.topography_dcmip200
+ClimaAtmos.topography_cosine_2d
+```
+
 ### Internals
 
 ```@docs

src/solver/type_getters.jl

@@ -429,36 +429,37 @@ function get_initial_condition(parsed_args, atmos)
     end
 end
 
+function get_topography(parsed_args)
+    topo_str = parsed_args["topography"]
+    topo_types = Dict("NoWarp" => NoTopography(),
+        "Cosine2D" => CosineTopography2D(),
+        "Cosine3D" => CosineTopography3D(),
+        "Agnesi" => AgnesiTopography(),
+        "Schar" => ScharTopography(),
+        "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(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)
+    # Main.@infiltrate
     steady_state_velocity =
         get_steady_state_velocity(params, Y, config.parsed_args)
 

src/topography/steady_state_solutions.jl

@@ -1,4 +1,8 @@
-import StaticArrays: @SVector
+import StaticArrays: SVector, @SVector
+# import .Topography: NoTopography, AgnesiTopography, ScharTopography, CosineTopography
+# import .Topography: topography_agnesi, topography_schar, topography_cosine
+# import .Topography: agnesi_params, schar_params, cosine_params
+
 
 background_u(::Type{FT}) where {FT} = FT(10)
 background_N(::Type{FT}) where {FT} = FT(0.01) # This needs to be a small value.
@@ -167,23 +171,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)
+    x = coord.x
+    z = coord.z
+    return steady_state_velocity_cosine(params, x, FT(0), z, λ, λ, z_top)
 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
+    x = coord.x
+    y = coord.y
+    z = coord.z
     return steady_state_velocity_cosine(params, x, y, z, λ, λ, z_top)
 end
 
@@ -253,7 +260,7 @@ 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(::AgnesiTopography, params, coord, z_top)
     FT = eltype(params)
     (; h_max, x_center, a) = agnesi_params(FT)
     topography_agnesi_Fh(k_x) = h_max * a / 2 * exp(-a * abs(k_x))
@@ -270,7 +277,7 @@ function steady_state_velocity_agnesi(params, coord, z_top)
     )
 end
 
-function steady_state_velocity_schar(params, coord, z_top)
+function steady_state_velocity(::ScharTopography, params, coord, z_top)
     FT = eltype(params)
     (; h_max, x_center, λ, a) = schar_params(FT)
     k_peak = 2 * FT(π) / λ

src/topography/topography.jl

@@ -1,3 +1,16 @@
+# module Topography
+
+using ClimaCore: Geometry, Spaces, Fields
+
+# export NoTopography, AgnesiTopography, ScharTopography, CosineTopography
+# export EarthTopography, DCMIP200Topography, Hughes2023Topography
+
+# export topography_dcmip200, topography_hughes2023
+# export topography_agnesi, agnesi_params
+# export topography_schar, schar_params
+# export topography_cosine_2d, topography_cosine_3d
+# export topography_cosine, cosine_params
+
 """
     topography_dcmip200(coord)
 
@@ -105,11 +118,109 @@ Surface elevation for 3D cosine hills.
 """
 function topography_cosine_3d(coord)
     FT = Geometry.float_type(coord)
-    (; x, y) = coord
-    (; h_max, λ) = cosine_params(FT)
+    x, y = coord.x, coord.y
+    cos_params = cosine_params(FT)
+    h_max, λ = cos_params.h_max, cos_params.λ
     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))
+
+abstract type AbstractTopography end
+
+topo_func(t::AbstractTopography) = error("No topography function for topography $t")
+
+struct NoTopography <: AbstractTopography end
+
+# Analytical topography types for idealized test cases
+
+"""
+    CosineTopography(; h_max = 25, λ = 25e3, dim = 2)
+
+Cosine hill topography in 2D or 3D.
+
+# Arguments
+- `h_max::FT`: Maximum elevation (m)
+- `λ::FT`: Wavelength of the cosine hills (m)
+- `dim::Int`: Spatial dimension (2 or 3)
+"""
+Base.@kwdef struct CosineTopography{D, FT} <: AbstractTopography
+    h_max::FT = 25.0
+    λ::FT = 25e3
+end
+
+CosineTopography2D() = CosineTopography{2, Float64}()
+CosineTopography3D() = CosineTopography{3, Float64}()
+
+topo_func(t::CosineTopography{2}) = topography_cosine_2d
+topo_func(t::CosineTopography{3}) = topography_cosine_3d
+
+"""
+    AgnesiTopography(; 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
+
+topo_func(::AgnesiTopography) = topography_agnesi
+
+
+"""
+    ScharTopography(; 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
+topo_func(::ScharTopography) = topography_schar
+
+
+# Data-based topography types
+
+"""
+    EarthTopography()
+
+Earth topography from ETOPO2022 data files.
+"""
+struct EarthTopography <: AbstractTopography end
+
+"""
+    DCMIP200Topography()
+
+Surface elevation for the DCMIP-2-0-0 test problem.
+"""
+struct DCMIP200Topography <: AbstractTopography end
+
+topo_func(::DCMIP200Topography) = topography_dcmip200
+
+"""
+    Hughes2023Topography()
+
+Surface elevation for baroclinic wave test from Hughes and Jablonowski (2023).
+"""
+struct Hughes2023Topography <: AbstractTopography end
+
+topo_func(::Hughes2023Topography) = topography_hughes2023
+
+# end

src/utils/common_spaces.jl

@@ -96,21 +96,11 @@ function make_hybrid_spaces(
     )
     z_grid = Grids.FiniteDifferenceGrid(z_topology)
 
-    topography = parsed_args["topography"]
-    @assert topography in (
-        "NoWarp",
-        "Earth",
-        "DCMIP200",
-        "Hughes2023",
-        "Agnesi",
-        "Schar",
-        "Cosine2D",
-        "Cosine3D",
-    )
-    if topography == "NoWarp"
+    topography = get_topography(parsed_args)
+    if topography isa NoTopography
         z_surface = zeros(h_space)
         @info "No surface orography warp applied"
-    elseif topography == "Earth"
+    elseif topography isa EarthTopography
         z_surface = SpaceVaryingInput(
             AA.earth_orography_file_path(;
                 context = ClimaComms.context(h_space),
@@ -120,26 +110,14 @@ function make_hybrid_spaces(
         )
         @info "Remapping Earth orography from ETOPO2022 data onto horizontal space"
     else
-        topography_function = if topography == "DCMIP200"
-            topography_dcmip200
-        elseif topography == "Hughes2023"
-            topography_hughes2023
-        elseif topography == "Agnesi"
-            topography_agnesi
-        elseif topography == "Schar"
-            topography_schar
-        elseif topography == "Cosine2D"
-            topography_cosine_2d
-        elseif topography == "Cosine3D"
-            topography_cosine_3d
-        end
+        topography_function = topo_func(topography)
         z_surface = SpaceVaryingInput(topography_function, h_space)
-        @info "Using $topography orography"
+        @info "Using $(nameof(typeof(topography))) orography"
     end
 
-    if topography == "NoWarp"
+    if topography isa NoTopography
         hypsography = Hypsography.Flat()
-    elseif topography == "Earth"
+    elseif topography isa EarthTopography
         mask(x::FT) where {FT} = x * FT(x > 0)
         z_surface = @. mask(z_surface)
         # diff_cfl = νΔt/Δx²