Add topography module and types

Author: nefrathenrici

docs/src/api.md

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

src/ClimaAtmos.jl

@@ -31,6 +31,7 @@ include(joinpath("utils", "AtmosArtifacts.jl"))
 import .AtmosArtifacts as AA
 
 include(joinpath("topography", "topography.jl"))
+include(joinpath("topography", "steady_state_velocity.jl"))
 include(joinpath("topography", "steady_state_solutions.jl"))
 
 include(joinpath("parameterized_tendencies", "radiation", "RRTMGPInterface.jl"))

src/solver/type_getters.jl

@@ -424,39 +424,30 @@ function get_initial_condition(parsed_args, atmos)
     end
 end
 
+function get_topography(parsed_args)
+    topo_str = parsed_args["topography"]
+    topo_types = Dict("NoWarp" => Topography.NoTopography(),
+        "Cosine2D" => Topography.CosineTopography(; dim = 2),
+        "Cosine3D" => Topography.CosineTopography(; dim = 3),
+        "Agnesi" => Topography.AgnesiTopography(),
+        "Schar" => Topography.ScharTopography(),
+        "Earth" => Topography.EarthTopography(),
+        "Hughes2023" => Topography.Hughes2023Topography(),
+        "DCMIP200" => Topography.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
-    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.(params, Fields.coordinate_field(Y.f), z_top)
-    end
-    @info "Steady-state velocity approximation completed: $s"
-    return (; ᶜu, ᶠu)
+    topography = get_topography(parsed_args)
+    return compute_steady_state_velocity(topography, params, Y)
 end
 
 function get_surface_setup(parsed_args)

src/topography/steady_state_velocity.jl

@@ -0,0 +1,58 @@
+import .Topography:
+    NoTopography, CosineTopography, AgnesiTopography, ScharTopography
+
+"""
+    compute_steady_state_velocity(topography, params, z_top)
+
+Compute the steady-state velocity for the given topography type.
+"""
+function compute_steady_state_velocity(::NoTopography, params, Y)
+    top_level = Spaces.nlevels(axes(Y.c)) + Fields.half
+    z_top = Fields.level(Fields.coordinate_field(Y.f).z, top_level)
+
+    ᶜu = steady_state_velocity_no_warp.(params, Fields.coordinate_field(Y.c), z_top)
+    ᶠu = steady_state_velocity_no_warp.(params, Fields.coordinate_field(Y.f), z_top)
+
+    return (; ᶜu, ᶠu)
+end
+
+function compute_steady_state_velocity(topography::CosineTopography, params, Y)
+    top_level = Spaces.nlevels(axes(Y.c)) + Fields.half
+    z_top = Fields.level(Fields.coordinate_field(Y.f).z, top_level)
+
+    if topography.dimension == 2
+        ᶜu = steady_state_velocity_cosine_2d.(params, Fields.coordinate_field(Y.c), z_top)
+        ᶠu = steady_state_velocity_cosine_2d.(params, Fields.coordinate_field(Y.f), z_top)
+    elseif topography.dimension == 3
+        ᶜu = steady_state_velocity_cosine_3d.(params, Fields.coordinate_field(Y.c), z_top)
+        ᶠu = steady_state_velocity_cosine_3d.(params, Fields.coordinate_field(Y.f), z_top)
+    else
+        throw(
+            ArgumentError(
+                "CosineTopography dimension must be 2 or 3, got $(topography.dimension)",
+            ),
+        )
+    end
+
+    return (; ᶜu, ᶠu)
+end
+
+function compute_steady_state_velocity(::AgnesiTopography, params, Y)
+    top_level = Spaces.nlevels(axes(Y.c)) + Fields.half
+    z_top = Fields.level(Fields.coordinate_field(Y.f).z, top_level)
+
+    ᶜu = steady_state_velocity_agnesi.(params, Fields.coordinate_field(Y.c), z_top)
+    ᶠu = steady_state_velocity_agnesi.(params, Fields.coordinate_field(Y.f), z_top)
+
+    return (; ᶜu, ᶠu)
+end
+
+function compute_steady_state_velocity(topography::ScharTopography, params, Y)
+    top_level = Spaces.nlevels(axes(Y.c)) + Fields.half
+    z_top = Fields.level(Fields.coordinate_field(Y.f).z, top_level)
+
+    ᶜu = steady_state_velocity_schar.(params, Fields.coordinate_field(Y.c), z_top)
+    ᶠu = steady_state_velocity_schar.(params, Fields.coordinate_field(Y.f), z_top)
+
+    return (; ᶜu, ᶠu)
+end

src/topography/topography.jl

@@ -1,3 +1,13 @@
+module Topography
+
+using ClimaCore: Geometry
+
+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)
 
@@ -113,3 +123,97 @@ 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 <: AbstractTopography
+    h_max = 25
+    λ = 25e3
+    dim::Int = 2
+end
+
+topo_func(t::CosineTopography) = t.dim == 2 ? topography_cosine_2d : 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 <: AbstractTopography
+    h_max = 25
+    x_center = 50e3
+    a = 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 <: AbstractTopography
+    h_max = 25
+    x_center = 50e3
+    λ = 4e3
+    a = 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 Topography.NoTopography
         z_surface = zeros(h_space)
         @info "No surface orography warp applied"
-    elseif topography == "Earth"
+    elseif topography isa Topography.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 = Topography.topo_func(topography)
         z_surface = SpaceVaryingInput(topography_function, h_space)
         @info "Using $topography orography"
     end
 
-    if topography == "NoWarp"
+    if topography isa Topography.NoTopography
         hypsography = Hypsography.Flat()
-    elseif topography == "Earth"
+    elseif topography isa Topography.EarthTopography
         mask(x::FT) where {FT} = x * FT(x > 0)
         z_surface = @. mask(z_surface)
         # diff_cfl = νΔt/Δx²

test/topography.jl

@@ -30,8 +30,8 @@ include("test_helpers.jl")
 
 @testset "test topography functions" begin
     (; FT, coords) = get_spherical_spaces()
-    @test extrema(CA.topography_dcmip200.(coords)) == (FT(0), FT(2000))
-    loc = findmax(parent(CA.topography_dcmip200.(coords)))
+    @test extrema(CA.Topography.topography_dcmip200.(coords)) == (FT(0), FT(2000))
+    loc = findmax(parent(CA.Topography.topography_dcmip200.(coords)))
     @test parent(coords.lat)[loc[2]] == FT(0)
     @test parent(coords.long)[loc[2]] == FT(-90)
 end