diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index 45d24d916f..d8195cdeba 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -1864,7 +1864,7 @@ version = "2.5.5+0" [[deps.OpenLibm_jll]] deps = ["Artifacts", "Libdl"] uuid = "05823500-19ac-5b8b-9628-191a04bc5112" -version = "0.8.5+0" +version = "0.8.1+2" [[deps.OpenMPI_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"] @@ -2210,9 +2210,9 @@ version = "0.5.1+0" [[deps.RootSolvers]] deps = ["ForwardDiff", "Printf"] -git-tree-sha1 = "75b130d104fb3d9d39a5c784fad23aebc7d7d956" +git-tree-sha1 = "769388dbf7656e70f6ee250f35bb6cbca8f43203" uuid = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74" -version = "1.0.0" +version = "0.4.6" [[deps.RoundingEmulator]] git-tree-sha1 = "40b9edad2e5287e05bd413a38f61a8ff55b9557b" diff --git a/.buildkite/PrecompileCI/src/PrecompileCI.jl b/.buildkite/PrecompileCI/src/PrecompileCI.jl index 23db0d43c9..3784ee3a2c 100644 --- a/.buildkite/PrecompileCI/src/PrecompileCI.jl +++ b/.buildkite/PrecompileCI/src/PrecompileCI.jl @@ -15,33 +15,39 @@ import ClimaParams x_max = y_max = 1e8 z_max = FT(30000.0) dz_bottom = FT(500) # other values? - z_stretch = Meshes.HyperbolicTangentStretching(dz_bottom) # Meshes.Uniform() + z_stretch = true # Meshes.Uniform() bubble = true # false - parsed_args = - Dict{String, Any}("topography" => "NoWarp", "topo_smoothing" => false) - comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded()) - deep = false - - # constants - quad = Quadratures.GLL{4}() + nh_poly = 3 # GLL{4} = nh_poly + 1 + context = ClimaComms.context(ClimaComms.CPUSingleThreaded()) + deep_atmosphere = false + topography = CA.NoTopography() params = CA.ClimaAtmosParameters(FT) radius = CA.Parameters.planet_radius(params) - # Sphere - horz_mesh = CA.cubed_sphere_mesh(; radius, h_elem) - h_space = CA.make_horizontal_space(horz_mesh, quad, comms_ctx, bubble) - CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args) - - # box - horizontal_mesh = CA.periodic_rectangle_mesh(; x_max, y_max, x_elem, y_elem) - h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) - # This is broken - # CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args) + sphere_grid = CA.SphereGrid( + FT; + context, radius, h_elem, nh_poly, z_elem, z_max, z_stretch, + dz_bottom, bubble, deep_atmosphere, + topography, + ) + box_grid = CA.BoxGrid( + FT; context, x_elem, x_max, y_elem, y_max, z_elem, z_max, nh_poly, + z_stretch, dz_bottom, bubble, deep_atmosphere, + periodic_x = true, periodic_y = true, topography, + ) + plane_grid = CA.PlaneGrid( + FT; context, x_elem, x_max, z_elem, z_max, nh_poly, z_stretch, + dz_bottom, bubble, deep_atmosphere, periodic_x = true, + topography, + ) + column_grid = CA.ColGrid( + FT; context, z_elem, z_max, z_stretch, dz_bottom, + ) - # plane - horizontal_mesh = CA.periodic_line_mesh(; x_max, x_elem) - h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) + for grid in [sphere_grid, box_grid, plane_grid, column_grid] + CA.get_spaces(grid, context) + end end end -end # module Precompile +end # module PrecompileCI diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 137edc9c07..410539d05f 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -28,8 +28,16 @@ steps: - julia --check-bounds=yes -e 'using TOML; @assert TOML.parsefile(".buildkite/Manifest-v1.11.toml")["deps"]["ClimaAtmos"][1]["path"] == ".."' - echo "--- Instantiate .buildkite" - - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true); using CUDA; CUDA.precompile_runtime(); Pkg.status()'" - + - julia --project=.buildkite -e ' + + using Pkg; + Pkg.Registry.update(); + Pkg.add(PackageSpec(name="ClimaCore", url="https://github.com/CliMA/ClimaCore.jl", rev="ne/fd_grid")); + Pkg.instantiate(); + Pkg.precompile(; strict=true); + using CUDA; CUDA.precompile_runtime(); + Pkg.status() + ' agents: slurm_cpus_per_task: 8 slurm_gpus: 1 diff --git a/Project.toml b/Project.toml index 7fbeff4497..f86a9c5cb7 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,7 @@ ArgParse = "1" Artifacts = "1" AtmosphericProfilesLibrary = "0.1.7" ClimaComms = "0.6.9" -ClimaCore = "0.14.37" +ClimaCore = "0.14.42" ClimaDiagnostics = "0.2.12" ClimaInterpolations = "0.1.0" ClimaParams = "1.0.2" diff --git a/calibration/experiments/gcm_driven_scm/helper_funcs.jl b/calibration/experiments/gcm_driven_scm/helper_funcs.jl index 6370559c92..c0ba82fe99 100644 --- a/calibration/experiments/gcm_driven_scm/helper_funcs.jl +++ b/calibration/experiments/gcm_driven_scm/helper_funcs.jl @@ -23,6 +23,7 @@ CLIMADIAGNOSTICS_LES_NAME_MAP = """Get z cell centers coordinates for CA run, given config. """ function get_z_grid(atmos_config; z_max = nothing) params = CA.ClimaAtmosParameters(atmos_config) + # TODO: Update this to use the new get_spaces function spaces = CA.get_spaces(atmos_config.parsed_args, params, atmos_config.comms_ctx) coord = CA.Fields.coordinate_field(spaces.center_space) diff --git a/docs/src/api.md b/docs/src/api.md index 64e9519f4f..bbdf601c9a 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -43,12 +43,21 @@ ClimaAtmos.InitialConditions.Bomex ClimaAtmos.InitialConditions.Soares ``` -### Helper +## Helper ```@docs ClimaAtmos.InitialConditions.ColumnInterpolatableField ``` +## Grids + +```@docs +ClimaAtmos.ColGrid + ClimaAtmos.SphereGrid +ClimaAtmos.PlaneGrid +ClimaAtmos.BoxGrid +``` + ## Jacobian ```@docs diff --git a/post_processing/ci_plots.jl b/post_processing/ci_plots.jl index 57fc0b45d1..ab62485f02 100644 --- a/post_processing/ci_plots.jl +++ b/post_processing/ci_plots.jl @@ -437,8 +437,6 @@ make_plots_generic( simulation_path, vars, time = LAST_SNAP, - x = 0.0, # Our columns are still 3D objects... - y = 0.0, more_kwargs = YLINEARSCALE, ) ``` @@ -453,8 +451,6 @@ make_plots_generic( simulation_path, vars, time = LAST_SNAP, - x = 0.0, # Our columns are still 3D objects... - y = 0.0, more_kwargs = YLINEARSCALE, ) ``` @@ -589,14 +585,23 @@ ColumnPlots = Union{ function make_plots(::ColumnPlots, output_paths::Vector{<:AbstractString}) simdirs = SimDir.(output_paths) short_names = ["ta", "wa"] - vars = map_comparison(get, simdirs, short_names) + vars = map_comparison(simdirs, short_names) do simdir, short_name + var = get(simdir; short_name) + # For vertical-only (FiniteDifferenceGrid) spaces, the data may have + # extra singleton dimensions. Check and squeeze if needed. + if haskey(var.dims, "x") && length(var.dims["x"]) == 1 + var = slice(var; x = var.dims["x"][1]) + end + if haskey(var.dims, "y") && length(var.dims["y"]) == 1 + var = slice(var; y = var.dims["y"][1]) + end + return var + end make_plots_generic( output_paths, vars, time = LAST_SNAP, - x = 0.0, # Our columns are still 3D objects... - y = 0.0, MAX_NUM_COLS = length(simdirs), more_kwargs = YLINEARSCALE, ) @@ -644,7 +649,7 @@ function make_plots( end vars = [ - slice(get(simdir; short_name), x = 0.0, y = 0.0) for + get(simdir; short_name) for short_name in short_names ] @@ -689,7 +694,7 @@ function make_plots( surface_precip = read_var(simdir.variable_paths["pr"]["inst"]["10s"]) viz.line_plot1D!( fig, - slice(surface_precip, x = 0.0, y = 0.0); + surface_precip; p_loc = [pr_row, 1:3], ) @@ -1371,6 +1376,12 @@ EDMFBoxPlots = Union{ Val{:diagnostic_edmfx_dycoms_rf01_box}, Val{:diagnostic_edmfx_trmm_box_0M}, Val{:diagnostic_edmfx_dycoms_rf01_explicit_box}, + Val{:prognostic_edmfx_bomex_box}, + Val{:rcemipii_box_diagnostic_edmfx}, + Val{:diagnostic_edmfx_trmm_stretched_box}, +} + +EDMFColumnPlots = Union{ Val{:prognostic_edmfx_adv_test_column}, Val{:prognostic_edmfx_gabls_column}, Val{:prognostic_edmfx_gabls_column_sparse_autodiff}, @@ -1386,13 +1397,10 @@ EDMFBoxPlots = Union{ Val{:prognostic_edmfx_simpleplume_column}, Val{:prognostic_edmfx_gcmdriven_column}, Val{:prognostic_edmfx_tv_era5driven_column}, - Val{:prognostic_edmfx_bomex_box}, - Val{:rcemipii_box_diagnostic_edmfx}, Val{:prognostic_edmfx_soares_column}, - Val{:diagnostic_edmfx_trmm_stretched_box}, } -EDMFBoxPlotsWithPrecip = Union{ +EDMFColumnPlotsWithPrecip = Union{ Val{:prognostic_edmfx_rico_column}, Val{:prognostic_edmfx_rico_implicit_column}, Val{:prognostic_edmfx_rico_column_2M}, @@ -1485,14 +1493,99 @@ end function make_plots( sim_type::Union{ EDMFBoxPlots, - EDMFBoxPlotsWithPrecip, DiagEDMFBoxPlotsWithPrecip, }, output_paths::Vector{<:AbstractString}, ) simdirs = SimDir.(output_paths) - if sim_type isa EDMFBoxPlotsWithPrecip + if sim_type isa DiagEDMFBoxPlotsWithPrecip + precip_names = ("husra", "hussn", "husraup", "hussnup") + else + precip_names = () + end + + short_names = [ + "wa", + "waup", + "ta", + "taup", + "hus", + "husup", + "arup", + "tke", + "ua", + "thetaa", + "thetaaup", + "ha", + "haup", + "hur", + "hurup", + "lmix", + "cl", + "clw", + "clwup", + "cli", + "cliup", + precip_names..., + ] + reduction = "inst" + + available_periods = ClimaAnalysis.available_periods( + simdirs[1]; + short_name = short_names[1], + reduction, + ) + # choose the shortest available period + available_periods = collect(available_periods) # ensure vector for indexing + period = available_periods[argmin(CA.time_to_seconds.(available_periods))] + + short_name_tuples = pair_edmf_names(short_names) + var_groups_zt = + map_comparison(simdirs, short_name_tuples) do simdir, name_tuple + return [ + slice( + get(simdir; short_name, reduction, period), + x = 0.0, + y = 0.0, + ) for short_name in name_tuple + ] + end + + var_groups_z = [ + ([slice(v, time = LAST_SNAP) for v in group]...,) for + group in var_groups_zt + ] + + tmp_file = make_plots_generic( + output_paths, + output_name = "tmp", + var_groups_z; + plot_fn = plot_edmf_vert_profile!, + MAX_NUM_COLS = 2, + MAX_NUM_ROWS = 4, + ) + + make_plots_generic( + output_paths, + vcat(var_groups_zt...), + plot_fn = plot_parsed_attribute_title!, + summary_files = [tmp_file], + MAX_NUM_COLS = 2, + MAX_NUM_ROWS = 4, + ) +end + +function make_plots( + sim_type::Union{ + EDMFColumnPlots, + EDMFColumnPlotsWithPrecip, + }, + output_paths::Vector{<:AbstractString}, +) + simdirs = SimDir.(output_paths) + + if sim_type isa EDMFColumnPlotsWithPrecip if sim_type isa Val{:prognostic_edmfx_rico_column_2M} precip_names = ( "husra", @@ -1512,8 +1605,6 @@ function make_plots( precip_names = ("husra", "hussn", "husraup", "hussnup", "husraen", "hussnen") end - elseif sim_type isa DiagEDMFBoxPlotsWithPrecip - precip_names = ("husra", "hussn", "husraup", "hussnup") else precip_names = () end @@ -1557,11 +1648,7 @@ function make_plots( var_groups_zt = map_comparison(simdirs, short_name_tuples) do simdir, name_tuple return [ - slice( - get(simdir; short_name, reduction, period), - x = 0.0, - y = 0.0, - ) for short_name in name_tuple + get(simdir; short_name, reduction, period) for short_name in name_tuple ] end diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index a7a8c587d1..5d42d4e45c 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -156,6 +156,7 @@ import .Diagnostics as CAD include(joinpath("callbacks", "get_callbacks.jl")) include(joinpath("simulation", "AtmosSimulations.jl")) +include(joinpath("simulation", "grids.jl")) include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters include(joinpath("solver", "type_getters.jl")) diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index b5cf4d432b..6dbed22083 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -28,7 +28,6 @@ function temporary_quantities(Y, atmos) center_space, face_space = axes(Y.c), axes(Y.f) FT = Spaces.undertype(center_space) - CTh = CTh_vector_type(Y.c) uvw_vec = UVW(FT(0), FT(0), FT(0)) return (; ᶠtemp_scalar = Fields.Field(FT, face_space), # ᶠp, ᶠρK_h @@ -81,7 +80,7 @@ function temporary_quantities(Y, atmos) # TODO: Remove this hack sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ # Implicit solver cache: - ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CTh{FT}}}), + ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CT12{FT}}}), ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}), ᶠp_grad_matrix = similar(Y.f, BidiagonalMatrixRow{C3{FT}}), ᶠbidiagonal_matrix_ct3 = similar(Y.f, BidiagonalMatrixRow{CT3{FT}}), diff --git a/src/callbacks/get_callbacks.jl b/src/callbacks/get_callbacks.jl index 2b92be2fb8..97b7b0ea8b 100644 --- a/src/callbacks/get_callbacks.jl +++ b/src/callbacks/get_callbacks.jl @@ -167,6 +167,7 @@ function get_diagnostics(parsed_args, atmos_model, Y, p, sim_info, output_dir) FT(time_to_seconds(parsed_args["t_end"]) - t_start), start_date; output_writer = netcdf_writer, + topography = has_topography(axes(Y.c)), )..., diagnostics..., ] diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl index 50f7f24b96..7125a9ac78 100644 --- a/src/diagnostics/default_diagnostics.jl +++ b/src/diagnostics/default_diagnostics.jl @@ -30,6 +30,7 @@ function default_diagnostics( duration, start_date::DateTime; output_writer, + topography = true, ) # Unfortunately, [] is not treated nicely in a map (we would like it to be "excluded"), # so we need to manually filter out the submodels that don't have defaults associated @@ -48,7 +49,7 @@ function default_diagnostics( # We use a map because we want to ensure that diagnostics is a well defined type, not # Any. This reduces latency. return vcat( - core_default_diagnostics(output_writer, duration, start_date), + core_default_diagnostics(output_writer, duration, start_date; topography), map(non_empty_fields) do field default_diagnostics( getfield(model, field), @@ -131,7 +132,7 @@ end ######## # Core # ######## -function core_default_diagnostics(output_writer, duration, start_date) +function core_default_diagnostics(output_writer, duration, start_date; topography = true) core_diagnostics = [ "ts", "ta", @@ -168,22 +169,31 @@ function core_default_diagnostics(output_writer, duration, start_date) min_func = (args...; kwargs...) -> hourly_min(FT, args...; kwargs...) max_func = (args...; kwargs...) -> hourly_max(FT, args...; kwargs...) end + core_diagnostics = if topography + [ + # We need to compute the topography at the beginning of the + # simulation (and only at the beginning), so we set + # output/compute_schedule_func to false. It is still computed at the + # very beginning + ScheduledDiagnostic(; + variable = get_diagnostic_variable("orog"), + output_schedule_func = (integrator) -> false, + compute_schedule_func = (integrator) -> false, + output_writer, + output_short_name = "orog_inst", + ), + average_func(core_diagnostics...; output_writer, start_date)..., + min_func("ts"; output_writer, start_date), + max_func("ts"; output_writer, start_date)] + else + [ + average_func(core_diagnostics...; output_writer, start_date)..., + min_func("ts"; output_writer, start_date), + max_func("ts"; output_writer, start_date), + ] + end - return [ - # We need to compute the topography at the beginning of the simulation (and only at - # the beginning), so we set output/compute_schedule_func to false. It is still - # computed at the very beginning - ScheduledDiagnostic(; - variable = get_diagnostic_variable("orog"), - output_schedule_func = (integrator) -> false, - compute_schedule_func = (integrator) -> false, - output_writer, - output_short_name = "orog_inst", - ), - average_func(core_diagnostics...; output_writer, start_date)..., - min_func("ts"; output_writer, start_date), - max_func("ts"; output_writer, start_date), - ] + return core_diagnostics end ################## diff --git a/src/parameterized_tendencies/sponge/viscous_sponge.jl b/src/parameterized_tendencies/sponge/viscous_sponge.jl index fcfc409849..7b61665eb0 100644 --- a/src/parameterized_tendencies/sponge/viscous_sponge.jl +++ b/src/parameterized_tendencies/sponge/viscous_sponge.jl @@ -13,7 +13,9 @@ import ClimaCore.Spaces as Spaces αₘ(s, z) * ζ_viscous(s, z, zmax) function viscous_sponge_tendency_uₕ(ᶜuₕ, s) - s isa Nothing && return NullBroadcasted() + if s isa Nothing || axes(ᶜuₕ) isa Spaces.FiniteDifferenceSpace + return NullBroadcasted() + end (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ)) zmax = z_max(axes(ᶠz)) return @. lazy( diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 3de2286a09..a09db27d67 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -31,7 +31,6 @@ end function hyperdiffusion_cache( Y, hyperdiff::ClimaHyperdiffusion, turbconv_model, moisture_model, microphysics_model, ) - quadrature_style = Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c))) FT = eltype(Y) n = n_mass_flux_subdomains(turbconv_model) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index cedbcba868..5a8c593c73 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -71,15 +71,14 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos) approximate_solve_iters, ) = alg FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) DiagonalRow = DiagonalMatrixRow{FT} TridiagonalRow = TridiagonalMatrixRow{FT} BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} - TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} + TridiagonalRow_ACT12 = TridiagonalMatrixRow{Adjoint{FT, CT12{FT}}} BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - BidiagonalRow_C3xACTh = - BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} + BidiagonalRow_C3xACT12 = + BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT12{FT})')} DiagonalRow_C3xACT3 = DiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} TridiagonalRow_C3xACT3 = @@ -150,7 +149,7 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos) MatrixFields.unrolled_map( name -> (name, @name(c.uₕ)) => - similar(Y.c, TridiagonalRow_ACTh), + similar(Y.c, TridiagonalRow_ACT12), active_scalar_names, ) : () )..., @@ -162,7 +161,7 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos) name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), active_scalar_names, )..., - (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh), + (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACT12), (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), ) @@ -396,7 +395,6 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) rs = p.atmos.rayleigh_sponge FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' cv_d = FT(CAP.cv_d(params)) @@ -442,10 +440,10 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) if use_derivative(topography_flag) @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow( - adjoint(CTh(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), + adjoint(CT12(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), ) else - @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ))) + @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CT12(ᶜuₕ))) end @. ∂ᶜK_∂ᶠu₃ = ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + diff --git a/src/simulation/grids.jl b/src/simulation/grids.jl new file mode 100644 index 0000000000..a350eee4eb --- /dev/null +++ b/src/simulation/grids.jl @@ -0,0 +1,393 @@ +using ClimaCore: Geometry, Hypsography, Fields, Spaces, Meshes +using ClimaCore.CommonGrids: + ExtrudedCubedSphereGrid, ColumnGrid, Box3DGrid, SliceXZGrid, DefaultZMesh +using ClimaUtilities: SpaceVaryingInputs.SpaceVaryingInput +import .AtmosArtifacts as AA +import ClimaComms + +export SphereGrid, ColGrid, BoxGrid, PlaneGrid + +abstract type MeshWarpType end + +struct LinearWarp <: MeshWarpType end +struct SLEVEWarp <: MeshWarpType end + +""" + SphereGrid(FT; kwargs...) + +Create an ExtrudedCubedSphereGrid with topography support. + +# Arguments +- `FT`: Floating point type + +# Keyword Arguments +- `context`: Communications context +- `z_elem::Int`: Number of vertical elements +- `z_max::Real`: Maximum height +- `z_stretch::Bool`: Whether to use vertical stretching +- `dz_bottom::Real`: Bottom layer thickness for stretching +- `radius::Real`: Earth radius +- `h_elem::Int`: Number of horizontal elements per panel +- `nh_poly::Int`: Polynomial order +- `bubble::Bool`: Enable bubble correction +- `deep_atmosphere::Bool`: Enable deep atmosphere +- `topography::AbstractTopography`: Topography type +- `topography_damping_factor::Real`: Topography damping +- `mesh_warp_type`: MeshWarpType: Mesh warping type ("SLEVE" or "Linear") +- `sleve_eta::Real`: SLEVE parameter +- `sleve_s::Real`: SLEVE parameter +- `topo_smoothing::Bool`: Apply topography smoothing +""" +function SphereGrid( + FT; + context = ClimaComms.context(), + z_elem::Int = 10, + z_max::Real = 30000.0, + z_stretch::Bool = true, + dz_bottom::Real = 500.0, + radius::Real = 6.371229e6, + h_elem::Int = 6, + nh_poly::Int = 3, + bubble::Bool = false, + deep_atmosphere::Bool = true, + topography::AbstractTopography = NoTopography(), + topography_damping_factor::Real = 5.0, + mesh_warp_type::MeshWarpType = SLEVEWarp(), + sleve_eta::Real = 0.7, + sleve_s::Real = 10.0, + topo_smoothing::Bool = false, +) + n_quad_points = nh_poly + 1 + + stretch = + z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform() + + hypsography_fun = hypsography_function_from_topography( + FT, topography, topography_damping_factor, mesh_warp_type, + sleve_eta, sleve_s, topo_smoothing, + ) + + global_geometry = if deep_atmosphere + Geometry.DeepSphericalGlobalGeometry{FT}(radius) + else + Geometry.ShallowSphericalGlobalGeometry{FT}(radius) + end + + grid = ExtrudedCubedSphereGrid( + FT; + z_elem, z_min = 0, z_max, radius, h_elem, + n_quad_points, + device = ClimaComms.device(context), + context, + stretch, + hypsography_fun, + global_geometry, + enable_bubble = bubble, + ) + + return grid +end + +""" + ColGrid(FT; kwargs...) + +Create a ColumnGrid. + +# Arguments +- `FT`: Floating point type + +# Keyword Arguments +- `context`: Communications context +- `z_elem::Int`: Number of vertical elements +- `z_max::Real`: Maximum height +- `z_stretch::Bool`: Whether to use vertical stretching +- `dz_bottom::Real`: Bottom layer thickness for stretching +""" +function ColGrid( + FT; + context = ClimaComms.context(), + z_elem::Int = 10, + z_max::Real = 30000.0, + z_stretch::Bool = true, + dz_bottom::Real = 500.0, +) + stretch = + z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform() + z_mesh = DefaultZMesh( + FT; + z_min = 0, + z_max, + z_elem, + stretch, + ) + grid = ColumnGrid( + FT; + z_elem, z_min = 0, z_max, z_mesh, + device = ClimaComms.device(context), + context, + stretch, + ) + + return grid +end + +""" + BoxGrid(FT; kwargs...) + +Create a Box3DGrid with topography support. + +# Arguments +- `FT`: Floating point type + +# Keyword Arguments +- `context`: Communications context +- `x_elem::Int`: Number of x elements +- `x_max::Real`: Maximum x coordinate +- `y_elem::Int`: Number of y elements +- `y_max::Real`: Maximum y coordinate +- `z_elem::Int`: Number of vertical elements +- `z_max::Real`: Maximum height +- `nh_poly::Int`: Polynomial order +- `z_stretch::Bool`: Whether to use vertical stretching +- `dz_bottom::Real`: Bottom layer thickness for stretching +- `bubble::Bool`: Enable bubble correction +- `deep_atmosphere::Bool`: Enable deep atmosphere +- `periodic_x::Bool`: Periodic in x direction +- `periodic_y::Bool`: Periodic in y direction +- `topography::AbstractTopography`: Topography type +- `topography_damping_factor::Real`: Topography damping +- `mesh_warp_type::String`: Mesh warping type ("SLEVE" or "Linear") +- `sleve_eta::Real`: SLEVE parameter +- `sleve_s::Real`: SLEVE parameter +- `topo_smoothing::Bool`: Apply topography smoothing +""" +function BoxGrid( + FT; + context = ClimaComms.context(), + x_elem::Int = 6, + x_max::Real = 300000.0, + y_elem::Int = 6, + y_max::Real = 300000.0, + z_elem::Int = 10, + z_max::Real = 30000.0, + nh_poly::Int = 3, + z_stretch::Bool = true, + dz_bottom::Real = 500.0, + bubble::Bool = false, + deep_atmosphere::Bool = true, + periodic_x::Bool = true, + periodic_y::Bool = true, + topography::AbstractTopography = NoTopography(), + topography_damping_factor::Real = 5.0, + mesh_warp_type::MeshWarpType = LinearWarp(), + sleve_eta::Real = 0.7, + sleve_s::Real = 10.0, + topo_smoothing::Bool = false, +) + n_quad_points = nh_poly + 1 + + stretch = + z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform() + + hypsography_fun = hypsography_function_from_topography( + FT, topography, topography_damping_factor, mesh_warp_type, + sleve_eta, sleve_s, topo_smoothing, + ) + z_mesh = DefaultZMesh( + FT; + z_min = 0, + z_max, + z_elem, + stretch, + ) + grid = Box3DGrid( + FT; + z_elem, x_min = 0, x_max, y_min = 0, y_max, z_min = 0, z_max, + periodic_x, periodic_y, n_quad_points, x_elem, y_elem, + device = ClimaComms.device(context), + context, + stretch, + hypsography_fun, + global_geometry = Geometry.CartesianGlobalGeometry(), + z_mesh, + enable_bubble = bubble, + ) + + return grid +end + +""" + PlaneGrid(FT; kwargs...) + +Create a SliceXZGrid with topography support. + +# Arguments +- `FT`: Floating point type + +# Keyword Arguments +- `context`: Communications context +- `x_elem::Int`: Number of x elements +- `x_max::Real`: Maximum x coordinate +- `z_elem::Int`: Number of vertical elements +- `z_max::Real`: Maximum height +- `nh_poly::Int`: Polynomial order +- `z_stretch::Bool`: Whether to use vertical stretching +- `dz_bottom::Real`: Bottom layer thickness for stretching +- `bubble::Bool`: Enable bubble correction +- `deep_atmosphere::Bool`: Enable deep atmosphere +- `periodic_x::Bool`: Periodic in x direction +- `topography::AbstractTopography`: Topography type +- `topography_damping_factor::Real`: Topography damping +- `mesh_warp_type::String`: Mesh warping type ("SLEVE" or "Linear") +- `sleve_eta::Real`: SLEVE parameter +- `sleve_s::Real`: SLEVE parameter +- `topo_smoothing::Bool`: Apply topography smoothing +""" +function PlaneGrid( + FT; + context = ClimaComms.context(), + x_elem::Int = 6, + x_max::Real = 300000.0, + z_elem::Int = 10, + z_max::Real = 30000.0, + nh_poly::Int = 3, + z_stretch::Bool = true, + dz_bottom::Real = 500.0, + bubble::Bool = false, + deep_atmosphere::Bool = true, + periodic_x::Bool = true, + topography::AbstractTopography = NoTopography(), + topography_damping_factor::Real = 5.0, + mesh_warp_type::MeshWarpType = LinearWarp(), + sleve_eta::Real = 0.7, + sleve_s::Real = 10.0, + topo_smoothing::Bool = false, +) + n_quad_points = nh_poly + 1 + + stretch = + z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform() + + hypsography_fun = hypsography_function_from_topography( + FT, topography, topography_damping_factor, mesh_warp_type, + sleve_eta, sleve_s, topo_smoothing, + ) + + z_mesh = DefaultZMesh( + FT; + z_min = 0, + z_max, + z_elem, + stretch, + ) + + grid = SliceXZGrid( + FT; + z_elem, x_elem, x_min = 0, x_max, z_min = 0, z_max, z_mesh, + periodic_x, + n_quad_points, + device = ClimaComms.device(context), + context, + stretch, + hypsography_fun, + global_geometry = Geometry.CartesianGlobalGeometry(), + ) + + return grid +end + +""" + hypsography_function_from_topography( + topography, topography_damping_factor, mesh_warp_type, + sleve_eta, sleve_s, topo_smoothing, context) + +Create a hypsography function that handles topography integration. +""" +function hypsography_function_from_topography( + FT::Type{<:AbstractFloat}, + topography::AbstractTopography, + topography_damping_factor::Real, + mesh_warp_type, + sleve_eta::Real, + sleve_s::Real, + topo_smoothing::Bool, +) + return function (h_grid, z_grid) + topography isa NoTopography && return Hypsography.Flat() + + # Create horizontal space to work with topography + h_space = if h_grid isa Grids.SpectralElementGrid1D + Spaces.SpectralElementSpace1D(h_grid) + elseif h_grid isa Grids.SpectralElementGrid2D + Spaces.SpectralElementSpace2D(h_grid) + else + error("Unsupported horizontal grid type $(typeof(h_grid))") + end + + if topography isa EarthTopography + context = ClimaComms.context(h_space) + z_surface = SpaceVaryingInput( + AA.earth_orography_file_path(; context), + "z", + h_space, + ) + @info "Remapping Earth orography from ETOPO2022 data onto horizontal space" + else + z_surface = SpaceVaryingInput(topography_function(topography), h_space) + @info "Using $(nameof(typeof(topography))) orography" + end + + if topography isa EarthTopography + mask(x::FT) where {FT} = x * FT(x > 0) + z_surface = @. mask(z_surface) + # diff_cfl = νΔt/Δx² + diff_courant = 0.05 # Arbitrary example value. + Δh_scale = Spaces.node_horizontal_length_scale(h_space) + κ = FT(diff_courant * Δh_scale^2) + maxiter = Int(round(log(topography_damping_factor) / diff_courant)) + Hypsography.diffuse_surface_elevation!( + z_surface; + κ, + dt = FT(1), + maxiter, + ) + # Coefficient for horizontal diffusion may alternatively be + # determined from the empirical parameters suggested by + # E3SM v1/v2 Topography documentation found here: + # https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/1456603764/V1+Topography+GLL+grids + z_surface = @. mask(z_surface) + if mesh_warp_type isa SLEVEWarp + @info "SLEVE mesh warp" + hypsography = Hypsography.SLEVEAdaption( + Geometry.ZPoint.(z_surface), + FT(sleve_eta), + FT(sleve_s), + ) + elseif mesh_warp_type isa LinearWarp + @info "Linear mesh warp" + hypsography = + Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface)) + else + error("Undefined mesh-warping option $(nameof(typeof(mesh_warp_type)))") + end + else + # DCMIP200Topography, Hughes2023Topography, etc. + if topo_smoothing + Hypsography.diffuse_surface_elevation!(z_surface) + end + if mesh_warp_type isa SLEVEWarp + @info "SLEVE mesh warp" + hypsography = Hypsography.SLEVEAdaption( + Geometry.ZPoint.(z_surface), + FT(sleve_eta), + FT(sleve_s), + ) + elseif mesh_warp_type isa LinearWarp + @info "Linear mesh warp" + hypsography = + Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface)) + end + end + return hypsography + end +end diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 25cb640466..8e76420082 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -6,7 +6,7 @@ import ClimaUtilities.OutputPathGenerator import ClimaCore: InputOutput, Meshes, Spaces, Quadratures import ClimaAtmos.RRTMGPInterface as RRTMGPI import ClimaAtmos as CA -import ClimaCore.Fields +import ClimaCore: Fields, Grids import ClimaTimeSteppers as CTS import Logging @@ -240,127 +240,44 @@ function get_numerics(parsed_args, FT) return numerics end -function get_spaces(parsed_args, params, comms_ctx) - - FT = eltype(params) - z_elem = Int(parsed_args["z_elem"]) - z_max = FT(parsed_args["z_max"]) - dz_bottom = FT(parsed_args["dz_bottom"]) - bubble = parsed_args["bubble"] - deep = parsed_args["deep_atmosphere"] - - h_elem = parsed_args["h_elem"] - radius = CAP.planet_radius(params) - center_space, face_space = if parsed_args["config"] == "sphere" - nh_poly = parsed_args["nh_poly"] - quad = Quadratures.GLL{nh_poly + 1}() - horizontal_mesh = cubed_sphere_mesh(; radius, h_elem) - h_space = - make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) - z_stretch = if parsed_args["z_stretch"] - Meshes.HyperbolicTangentStretching(dz_bottom) - else - Meshes.Uniform() - end - make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; deep, parsed_args) - elseif parsed_args["config"] == "column" # single column - @warn "perturb_initstate flag is ignored for single column configuration" - FT = eltype(params) - Δx = FT(1) # Note: This value shouldn't matter, since we only have 1 column. - quad = Quadratures.GL{1}() - horizontal_mesh = periodic_rectangle_mesh(; - x_max = Δx, - y_max = Δx, - x_elem = 1, - y_elem = 1, - ) - if bubble - @warn "Bubble correction not compatible with single column configuration. It will be switched off." - bubble = false - end - h_space = - make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) - z_stretch = if parsed_args["z_stretch"] - Meshes.HyperbolicTangentStretching(dz_bottom) - else - Meshes.Uniform() - end - make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args) - elseif parsed_args["config"] == "box" - FT = eltype(params) - nh_poly = parsed_args["nh_poly"] - quad = Quadratures.GLL{nh_poly + 1}() - x_elem = Int(parsed_args["x_elem"]) - x_max = FT(parsed_args["x_max"]) - y_elem = Int(parsed_args["y_elem"]) - y_max = FT(parsed_args["y_max"]) - horizontal_mesh = periodic_rectangle_mesh(; - x_max = x_max, - y_max = y_max, - x_elem = x_elem, - y_elem = y_elem, - ) - h_space = - make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) - z_stretch = if parsed_args["z_stretch"] - Meshes.HyperbolicTangentStretching(dz_bottom) - else - Meshes.Uniform() - end - make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args, deep) - elseif parsed_args["config"] == "plane" - FT = eltype(params) - nh_poly = parsed_args["nh_poly"] - quad = Quadratures.GLL{nh_poly + 1}() - x_elem = Int(parsed_args["x_elem"]) - x_max = FT(parsed_args["x_max"]) - horizontal_mesh = - periodic_line_mesh(; x_max = x_max, x_elem = x_elem) - h_space = - make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble) - z_stretch = if parsed_args["z_stretch"] - Meshes.HyperbolicTangentStretching(dz_bottom) - else - Meshes.Uniform() - end - make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args, deep) - end - ncols = Fields.ncolumns(center_space) - ndofs_total = ncols * z_elem - hspace = Spaces.horizontal_space(center_space) - quad_style = Spaces.quadrature_style(hspace) - Nq = Quadratures.degrees_of_freedom(quad_style) - - @info "Resolution stats: " Nq h_elem z_elem ncols ndofs_total - return (; - center_space, - face_space, - horizontal_mesh, - quad, - z_max, - z_elem, - z_stretch, - ) -end - function get_spaces_restart(Y) center_space = axes(Y.c) face_space = axes(Y.f) return (; center_space, face_space) end +""" + get_spaces(grid, comms_ctx) + +Create center and face spaces from a ClimaCore grid. +""" +function get_spaces(grid, comms_ctx) + if grid isa Grids.ExtrudedFiniteDifferenceGrid + center_space = Spaces.CenterExtrudedFiniteDifferenceSpace(grid) + face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(grid) + elseif grid isa Grids.FiniteDifferenceGrid + center_space = Spaces.CenterFiniteDifferenceSpace(grid) + face_space = Spaces.FaceFiniteDifferenceSpace(grid) + else + error( + """Unsupported grid type: $(typeof(grid)). Expected \ + ExtrudedFiniteDifferenceGrid or FiniteDifferenceGrid""", + ) + end + return (; center_space, face_space) +end + function get_state_restart(config::AtmosConfig, restart_file, atmos_model_hash) (; parsed_args, comms_ctx) = config - sim_info = get_sim_info(config) + (; start_date) = get_sim_info(config) + use_itime = parsed_args["use_itime"] @assert !isnothing(restart_file) reader = InputOutput.HDF5Reader(restart_file, comms_ctx) Y = InputOutput.read_field(reader, "Y") # TODO: Do not use InputOutput.HDF5 directly t_start = InputOutput.HDF5.read_attribute(reader.file, "time") - t_start = - parsed_args["use_itime"] ? ITime(t_start; epoch = sim_info.start_date) : - t_start + t_start = use_itime ? ITime(t_start; epoch = start_date) : t_start if "atmos_model_hash" in keys(InputOutput.HDF5.attrs(reader.file)) atmos_model_hash_in_restart = InputOutput.HDF5.read_attribute(reader.file, "atmos_model_hash") @@ -627,39 +544,71 @@ function auto_detect_restart_file( return restart_file end -function get_sim_info(config::AtmosConfig) - (; comms_ctx, parsed_args) = config - FT = eltype(config) - (; job_id) = config +import ClimaUtilities.OutputPathGenerator + +""" + setup_output_dir(job_id, output_dir, output_dir_style, detect_restart_file, restart_file, comms_ctx) + +Unified function for setting up output directories and detecting restart files. +Used by both AtmosSimulation constructor and get_simulation. + +Returns a named tuple with: +- `output_dir`: The final output directory path +- `restart_file`: The restart file path (if any) +""" +function setup_output_dir( + job_id, + output_dir, + output_dir_style, + detect_restart_file, + restart_file, + comms_ctx, +) + # Set up base output directory default_output = haskey(ENV, "CI") ? job_id : joinpath("output", job_id) - out_dir = parsed_args["output_dir"] - base_output_dir = isnothing(out_dir) ? default_output : out_dir + base_output_dir = isnothing(output_dir) ? default_output : output_dir allowed_dir_styles = Dict( "activelink" => OutputPathGenerator.ActiveLinkStyle(), "removepreexisting" => OutputPathGenerator.RemovePreexistingStyle(), ) - requested_style = parsed_args["output_dir_style"] - - haskey(allowed_dir_styles, lowercase(requested_style)) || - error("output_dir_style $(requested_style) not available") + haskey(allowed_dir_styles, lowercase(output_dir_style)) || + error("output_dir_style $(output_dir_style) not available") - output_dir_style = allowed_dir_styles[lowercase(requested_style)] + output_dir_style_obj = allowed_dir_styles[lowercase(output_dir_style)] - # We look for a restart before creating a new output dir because we want to - # look for previous folders - restart_file = - parsed_args["detect_restart_file"] ? - auto_detect_restart_file(output_dir_style, base_output_dir) : - parsed_args["restart_file"] + final_restart_file = if detect_restart_file && isnothing(restart_file) + auto_detect_restart_file(output_dir_style_obj, base_output_dir) + else + restart_file + end output_dir = OutputPathGenerator.generate_output_path( base_output_dir; context = comms_ctx, - style = output_dir_style, + style = output_dir_style_obj, + ) + + return output_dir, final_restart_file +end + +function get_sim_info(config::AtmosConfig) + (; comms_ctx, parsed_args) = config + FT = eltype(config) + + (; job_id) = config + + output_dir, restart_file = CA.setup_output_dir( + job_id, + parsed_args["output_dir"], + parsed_args["output_dir_style"], + parsed_args["detect_restart_file"], + parsed_args["restart_file"], + comms_ctx, ) + if parsed_args["log_to_file"] @info "Logging to $output_dir/output.log" logger = ClimaComms.FileLogger(comms_ctx, output_dir) @@ -770,10 +719,102 @@ function get_comms_context(parsed_args) return comms_ctx end +function get_mesh_warp_type(s) + if s == "SLEVE" + return SLEVEWarp() + elseif s == "Linear" + return LinearWarp() + else + error("Unknown mesh warp type string: $s. Supported types are 'SLEVE' and 'Linear'") + end +end + +function get_grid(parsed_args, params, context) + FT = eltype(params) + if parsed_args["config"] == "sphere" + SphereGrid( + FT; + context, + radius = CAP.planet_radius(params), + h_elem = parsed_args["h_elem"], + nh_poly = parsed_args["nh_poly"], + z_elem = parsed_args["z_elem"], + z_max = parsed_args["z_max"], + z_stretch = parsed_args["z_stretch"], + dz_bottom = parsed_args["dz_bottom"], + bubble = parsed_args["bubble"], + deep_atmosphere = parsed_args["deep_atmosphere"], + topography = get_topography(FT, parsed_args), + topography_damping_factor = parsed_args["topography_damping_factor"], + mesh_warp_type = get_mesh_warp_type(parsed_args["mesh_warp_type"]), + sleve_eta = parsed_args["sleve_eta"], + sleve_s = parsed_args["sleve_s"], + topo_smoothing = parsed_args["topo_smoothing"], + ) + elseif parsed_args["config"] == "column" + ColGrid( + FT; + context, + z_elem = parsed_args["z_elem"], + z_max = parsed_args["z_max"], + z_stretch = parsed_args["z_stretch"], + dz_bottom = parsed_args["dz_bottom"], + ) + elseif parsed_args["config"] == "box" + BoxGrid( + FT; + context, + x_elem = parsed_args["x_elem"], + x_max = parsed_args["x_max"], + y_elem = parsed_args["y_elem"], + y_max = parsed_args["y_max"], + z_elem = parsed_args["z_elem"], + z_max = parsed_args["z_max"], + nh_poly = parsed_args["nh_poly"], + z_stretch = parsed_args["z_stretch"], + dz_bottom = parsed_args["dz_bottom"], + bubble = parsed_args["bubble"], + deep_atmosphere = parsed_args["deep_atmosphere"], + periodic_x = true, + periodic_y = true, + topography = get_topography(FT, parsed_args), + topography_damping_factor = parsed_args["topography_damping_factor"], + mesh_warp_type = get_mesh_warp_type(parsed_args["mesh_warp_type"]), + sleve_eta = parsed_args["sleve_eta"], + sleve_s = parsed_args["sleve_s"], + topo_smoothing = parsed_args["topo_smoothing"], + ) + elseif parsed_args["config"] == "plane" + PlaneGrid( + FT; + context, + x_elem = parsed_args["x_elem"], + x_max = parsed_args["x_max"], + z_elem = parsed_args["z_elem"], + z_max = parsed_args["z_max"], + nh_poly = parsed_args["nh_poly"], + z_stretch = parsed_args["z_stretch"], + dz_bottom = parsed_args["dz_bottom"], + bubble = parsed_args["bubble"], + deep_atmosphere = parsed_args["deep_atmosphere"], + periodic_x = true, + topography = get_topography(FT, parsed_args), + topography_damping_factor = parsed_args["topography_damping_factor"], + mesh_warp_type = get_mesh_warp_type(parsed_args["mesh_warp_type"]), + sleve_eta = parsed_args["sleve_eta"], + sleve_s = parsed_args["sleve_s"], + topo_smoothing = parsed_args["topo_smoothing"], + ) + end +end + function get_simulation(config::AtmosConfig) sim_info = get_sim_info(config) params = ClimaAtmosParameters(config) atmos = get_atmos(config, params) + comms_ctx = get_comms_context(config.parsed_args) + grid = get_grid(config.parsed_args, params, comms_ctx) + job_id = sim_info.job_id output_dir = sim_info.output_dir @info "Simulation info" job_id output_dir @@ -798,8 +839,10 @@ function get_simulation(config::AtmosConfig) end @info "Allocating Y: $s" else - spaces = get_spaces(config.parsed_args, params, config.comms_ctx) + spaces = get_spaces(grid, config.comms_ctx) end + @info spaces.center_space.grid + # TODO: add more information about the grid - stretch, etc. initial_condition = get_initial_condition(config.parsed_args, atmos) surface_setup = get_surface_setup(config.parsed_args) if !sim_info.restart @@ -886,7 +929,10 @@ function get_simulation(config::AtmosConfig) accum_str = join(CA.promote_period.(collect(periods_reductions)), ", ") checkpt_str = CA.promote_period(checkpoint_frequency) - @warn "The checkpointing frequency (dt_save_state_to_disk = $checkpt_str) should be an integer multiple of all diagnostics accumulation periods ($accum_str) so simulations can be safely restarted from any checkpoint" + @warn """The checkpointing frequency \ + (dt_save_state_to_disk = $checkpt_str) should be an integer \ + multiple of all diagnostics accumulation periods ($accum_str) \ + so simulations can be safely restarted from any checkpoint""" end end else @@ -910,11 +956,7 @@ function get_simulation(config::AtmosConfig) s = @timed_str begin integrator_args, integrator_kwargs = args_integrator( config.parsed_args, - Y, - p, - tspan, - ode_algo, - all_callbacks, + Y, p, tspan, ode_algo, all_callbacks, ) end diff --git a/src/solver/types.jl b/src/solver/types.jl index decb1efa7c..d02a2f4cba 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -1,3 +1,4 @@ +import ClimaComms import ClimaCore.Quadratures.GaussQuadrature as GQ import StaticArrays as SA import Thermodynamics as TD diff --git a/src/surface_conditions/surface_conditions.jl b/src/surface_conditions/surface_conditions.jl index 718709d7ed..9915ad3847 100644 --- a/src/surface_conditions/surface_conditions.jl +++ b/src/surface_conditions/surface_conditions.jl @@ -381,11 +381,11 @@ end #For non-RCEMIPII box models with prescribed surface temp, assume that the latitude is 0. function surface_temperature( ::Union{ZonallySymmetricSST, ZonallyAsymmetricSST}, - coordinates::Union{Geometry.XZPoint, Geometry.XYZPoint}, + coordinates::Union{Geometry.XZPoint, Geometry.XYZPoint, Geometry.ZPoint}, surface_temp_params, ) - (; x) = coordinates - FT = eltype(x) + (; z) = coordinates + FT = eltype(z) return FT(300) end diff --git a/src/utils/common_spaces.jl b/src/utils/common_spaces.jl index ccca9ef264..592040cc65 100644 --- a/src/utils/common_spaces.jl +++ b/src/utils/common_spaces.jl @@ -3,30 +3,6 @@ using ClimaCore: using ClimaComms using ClimaUtilities: SpaceVaryingInputs.SpaceVaryingInput -function periodic_line_mesh(; x_max, x_elem) - domain = Domains.IntervalDomain( - Geometry.XPoint(zero(x_max)), - Geometry.XPoint(x_max); - periodic = true, - ) - return Meshes.IntervalMesh(domain; nelems = x_elem) -end - -function periodic_rectangle_mesh(; x_max, y_max, x_elem, y_elem) - x_domain = Domains.IntervalDomain( - Geometry.XPoint(zero(x_max)), - Geometry.XPoint(x_max); - periodic = true, - ) - y_domain = Domains.IntervalDomain( - Geometry.YPoint(zero(y_max)), - Geometry.YPoint(y_max); - periodic = true, - ) - domain = Domains.RectangleDomain(x_domain, y_domain) - return Meshes.RectilinearMesh(domain, x_elem, y_elem) -end - # h_elem is the number of elements per side of every panel (6 panels in total) function cubed_sphere_mesh(; radius, h_elem) domain = Domains.SphereDomain(radius) @@ -71,105 +47,3 @@ function make_horizontal_space(mesh, quad, comms_ctx, bubble) end return space end - -function make_hybrid_spaces( - h_space, - z_max, - z_elem, - z_stretch; - deep = false, - parsed_args = nothing, -) - FT = eltype(z_max) - h_grid = Spaces.grid(h_space) - z_domain = Domains.IntervalDomain( - Geometry.ZPoint(zero(z_max)), - Geometry.ZPoint(z_max); - boundary_names = (:bottom, :top), - ) - z_mesh = Meshes.IntervalMesh(z_domain, z_stretch; nelems = z_elem) - @info "z heights" z_mesh.faces - device = ClimaComms.device(h_space) - z_topology = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(device), - z_mesh, - ) - z_grid = Grids.FiniteDifferenceGrid(z_topology) - - topography = get_topography(FT, parsed_args) - if topography isa NoTopography - z_surface = zeros(h_space) - @info "No surface orography warp applied" - elseif topography isa EarthTopography - z_surface = SpaceVaryingInput( - AA.earth_orography_file_path(; - context = ClimaComms.context(h_space), - ), - "z", - h_space, - ) - @info "Remapping Earth orography from ETOPO2022 data onto horizontal space" - else - z_surface = SpaceVaryingInput(topography_function(topography), h_space) - @info "Using $(nameof(typeof(topography))) orography" - end - - if topography isa NoTopography - hypsography = Hypsography.Flat() - elseif topography isa EarthTopography - mask(x::FT) where {FT} = x * FT(x > 0) - z_surface = @. mask(z_surface) - # diff_cfl = νΔt/Δx² - diff_courant = 0.05 # Arbitrary example value. - Δh_scale = Spaces.node_horizontal_length_scale(h_space) - κ = FT(diff_courant * Δh_scale^2) - n_attenuation = parsed_args["topography_damping_factor"] - maxiter = Int(round(log(n_attenuation) / diff_courant)) - Hypsography.diffuse_surface_elevation!( - z_surface; - κ, - dt = FT(1), - maxiter, - ) - # Coefficient for horizontal diffusion may alternatively be - # determined from the empirical parameters suggested by - # E3SM v1/v2 Topography documentation found here: - # https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/1456603764/V1+Topography+GLL+grids - z_surface = @. mask(z_surface) - if parsed_args["mesh_warp_type"] == "SLEVE" - @info "SLEVE mesh warp" - hypsography = Hypsography.SLEVEAdaption( - Geometry.ZPoint.(z_surface), - FT(parsed_args["sleve_eta"]), - FT(parsed_args["sleve_s"]), - ) - elseif parsed_args["mesh_warp_type"] == "Linear" - @info "Linear mesh warp" - hypsography = - Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface)) - else - @error "Undefined mesh-warping option" - end - else - if parsed_args["topo_smoothing"] - Hypsography.diffuse_surface_elevation!(z_surface) - end - if parsed_args["mesh_warp_type"] == "SLEVE" - @info "SLEVE mesh warp" - hypsography = Hypsography.SLEVEAdaption( - Geometry.ZPoint.(z_surface), - FT(parsed_args["sleve_eta"]), - FT(parsed_args["sleve_s"]), - ) - elseif parsed_args["mesh_warp_type"] == "Linear" - @info "Linear mesh warp" - hypsography = - Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface)) - end - end - - grid = Grids.ExtrudedFiniteDifferenceGrid(h_grid, z_grid, hypsography; deep) - center_space = Spaces.CenterExtrudedFiniteDifferenceSpace(grid) - face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(grid) - return center_space, face_space -end diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index fe0b74a10a..bae035c908 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -5,6 +5,7 @@ import ClimaComms import ClimaCore: Spaces, Topologies, Fields, Geometry import LinearAlgebra: norm_sqr using Dates: DateTime, @dateformat_str +import StaticArrays: SVector, SMatrix is_energy_var(symbol) = symbol in (:ρe_tot, :ρae_tot) is_momentum_var(symbol) = symbol in (:uₕ, :ρuₕ, :u₃, :ρw) @@ -222,41 +223,32 @@ Extracts the `g³ʰ` sub-tensor from the `gⁱʲ` tensor. """ function g³ʰ(gⁱʲ) full_CT_axis = axes(gⁱʲ)[1] - CTh_axis = if full_CT_axis == Geometry.Contravariant123Axis() - Geometry.Contravariant12Axis() - elseif full_CT_axis == Geometry.Contravariant13Axis() - Geometry.Contravariant1Axis() - elseif full_CT_axis == Geometry.Contravariant23Axis() - Geometry.Contravariant2Axis() - else - error("$full_CT_axis is missing either vertical or horizontal sub-axes") - end N = length(full_CT_axis) - return Geometry.AxisTensor( - (Geometry.Contravariant3Axis(), CTh_axis), - view(Geometry.components(gⁱʲ), N:N, 1:(N - 1)), - ) -end - -""" - CTh_vector_type(space) - -Extracts the (abstract) horizontal contravariant vector type from the given -`AbstractSpace`. -""" -function CTh_vector_type(space) - full_CT_axis = axes(eltype(Fields.local_geometry_field(space).gⁱʲ))[1] - return if full_CT_axis == Geometry.Contravariant123Axis() - Geometry.Contravariant12Vector + gⁱʲ_components = Geometry.components(gⁱʲ) + FT = eltype(gⁱʲ_components) + g³ʰ_components = if full_CT_axis == Geometry.Contravariant123Axis() + @inbounds SMatrix{1, 2, FT, 2}( + gⁱʲ_components[N, 1], + gⁱʲ_components[N, 2], + ) elseif full_CT_axis == Geometry.Contravariant13Axis() - Geometry.Contravariant1Vector + @inbounds val = gⁱʲ_components[N, 1] + SMatrix{1, 2, FT, 2}(val, zero(FT)) elseif full_CT_axis == Geometry.Contravariant23Axis() - Geometry.Contravariant2Vector + @inbounds val = gⁱʲ_components[N, 1] + SMatrix{1, 2, FT, 2}(zero(FT), val) else error("$full_CT_axis is missing either vertical or horizontal sub-axes") end + # Construct AxisTensor - the constructor at axistensors.jl:116-120 should be selected + # when components is already a StaticArray, avoiding the fallback that uses map/apply_type + axes_tuple = (Geometry.Contravariant3Axis(), Geometry.Contravariant12Axis()) + # Force selection of the specific constructor by ensuring types match exactly + # This avoids the fallback constructor at line 122 that uses map/apply_type + return Geometry.AxisTensor(axes_tuple, g³ʰ_components) end +has_topography(space::Spaces.FiniteDifferenceSpace) = false has_topography(space) = Spaces.grid(space).hypsography != Spaces.Grids.Flat() """ @@ -351,6 +343,9 @@ function do_dss(space::Spaces.AbstractSpace) Quadratures.GLL end +function do_dss(space::Spaces.FiniteDifferenceSpace) + return false +end using ClimaComms is_distributed(::ClimaComms.SingletonCommsContext) = false @@ -582,18 +577,8 @@ function parse_date(date_str) ) end -function iscolumn(space) - # TODO: Our columns are 2+1D boxes with one element at the base. Fix this - isbox = - Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))) isa - Domains.RectangleDomain - isbox || return false - has_one_element = - Meshes.nelements( - Spaces.topology(Spaces.horizontal_space(space)).mesh, - ) == 1 - has_one_element && return true -end +iscolumn(space::Spaces.FiniteDifferenceSpace) = true +iscolumn(space) = false function issphere(space) return Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))) isa diff --git a/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_3d.jl b/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_3d.jl index 14504e653a..a74aed0c98 100644 --- a/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_3d.jl +++ b/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_3d.jl @@ -16,14 +16,7 @@ include( ) include("../gw_plotutils.jl") -comms_ctx = ClimaComms.SingletonCommsContext() -(; config_file, job_id) = CA.commandline_kwargs() -config = CA.AtmosConfig(config_file; job_id, comms_ctx) - -config.parsed_args["topography"] = "Earth"; -config.parsed_args["topo_smoothing"] = false; -config.parsed_args["mesh_warp_type"] = "Linear"; -(; parsed_args) = config +context = ClimaComms.SingletonCommsContext() # load gfdl data include(joinpath(@__DIR__, "../../../artifact_funcs.jl")) @@ -86,12 +79,23 @@ z_elem = 33 dz_bottom = 300.0 radius = 6.371229e6 -quad = Quadratures.GLL{nh_poly + 1}() -horizontal_mesh = CA.cubed_sphere_mesh(; radius, h_elem) -h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, false) -z_stretch = Meshes.HyperbolicTangentStretching(dz_bottom) -center_space, face_space = - CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args) +grid = CA.SphereGrid( + FT; + context, + z_elem, + z_max, + z_stretch = true, + dz_bottom, + radius, + h_elem, + nh_poly, + bubble = false, + topography = CA.EarthTopography(), + topography_damping_factor = 5, + mesh_warp_type = CA.LinearWarp(), + topo_smoothing = false, +) +(; center_space, face_space) = CA.get_spaces(grid, context) ᶜlocal_geometry = Fields.local_geometry_field(center_space) ᶠlocal_geometry = Fields.local_geometry_field(face_space) diff --git a/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_baseflux.jl b/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_baseflux.jl index 0196e7e619..ad0fbd36c0 100644 --- a/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_baseflux.jl +++ b/test/parameterized_tendencies/gravity_wave/orographic_gravity_wave/ogwd_baseflux.jl @@ -12,10 +12,7 @@ include( joinpath(pkgdir(ClimaAtmos), "post_processing/remap", "remap_helpers.jl"), ) -comms_ctx = ClimaComms.SingletonCommsContext() -(; config_file, job_id) = CA.commandline_kwargs() -config = CA.AtmosConfig(config_file; job_id, comms_ctx) -config.parsed_args["topography"] = "NoWarp" +context = ClimaComms.SingletonCommsContext() # Create meshes and spaces h_elem = 6 @@ -24,18 +21,20 @@ z_max = 30e3 z_elem = 1 radius = 6.371229e6 -quad = Quadratures.GLL{nh_poly + 1}() -horizontal_mesh = CA.cubed_sphere_mesh(; radius, h_elem) -h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, false) -z_stretch = Meshes.Uniform() -center_space, face_space = CA.make_hybrid_spaces( - h_space, - z_max, +grid = CA.SphereGrid( + FT; + context, z_elem, - z_stretch; - parsed_args = config.parsed_args, + z_max, + z_stretch = false, + radius, + h_elem, + nh_poly, + bubble = false, + topography = CA.NoTopography(), ) - +(; center_space, face_space) = CA.get_spaces(grid, context) +h_space = Spaces.horizontal_space(center_space) ᶜlocal_geometry = Fields.local_geometry_field(center_space) ᶠlocal_geometry = Fields.local_geometry_field(face_space) diff --git a/test/test_helpers.jl b/test/test_helpers.jl index 10092d87b0..a72a08b2e1 100644 --- a/test/test_helpers.jl +++ b/test/test_helpers.jl @@ -28,6 +28,30 @@ function generate_test_simulation(config) return (; Y = Y, p = p, params = p.params, simulation = simulation) end +function periodic_line_mesh(; x_max, x_elem) + domain = Domains.IntervalDomain( + Geometry.XPoint(zero(x_max)), + Geometry.XPoint(x_max); + periodic = true, + ) + return Meshes.IntervalMesh(domain; nelems = x_elem) +end + +function periodic_rectangle_mesh(; x_max, y_max, x_elem, y_elem) + x_domain = Domains.IntervalDomain( + Geometry.XPoint(zero(x_max)), + Geometry.XPoint(x_max); + periodic = true, + ) + y_domain = Domains.IntervalDomain( + Geometry.YPoint(zero(y_max)), + Geometry.YPoint(y_max); + periodic = true, + ) + domain = Domains.RectangleDomain(x_domain, y_domain) + return Meshes.RectilinearMesh(domain, x_elem, y_elem) +end + function get_spherical_spaces(; FT = Float32) context = ClimaComms.SingletonCommsContext() radius = FT(10π) diff --git a/test/utilities.jl b/test/utilities.jl index 27b35cecf4..82edbfa2b3 100644 --- a/test/utilities.jl +++ b/test/utilities.jl @@ -298,7 +298,7 @@ end @testset "interval domain" begin # Interval Spaces (; zlim, velem) = get_cartesian_spaces() - line_mesh = CA.periodic_line_mesh(; x_max = zlim[2], x_elem = velem) + line_mesh = periodic_line_mesh(; x_max = zlim[2], x_elem = velem) @test line_mesh isa Meshes.IntervalMesh @test Geometry.XPoint(zlim[1]) == Meshes.domain(line_mesh).coord_min @test Geometry.XPoint(zlim[2]) == Meshes.domain(line_mesh).coord_max @@ -308,7 +308,7 @@ end @testset "periodic rectangle meshes (spectral elements)" begin # Interval Spaces (; xlim, zlim, velem, helem, npoly) = get_cartesian_spaces() - rectangle_mesh = CA.periodic_rectangle_mesh(; + rectangle_mesh = periodic_rectangle_mesh(; x_max = xlim[2], y_max = xlim[2], x_elem = helem, @@ -329,14 +329,14 @@ end comms_ctx = ClimaComms.context(device) FT = eltype(xlim) # 1D Space - line_mesh = CA.periodic_line_mesh(; x_max = zlim[2], x_elem = velem) + line_mesh = periodic_line_mesh(; x_max = zlim[2], x_elem = velem) @test line_mesh isa Meshes.AbstractMesh1D horz_plane_space = CA.make_horizontal_space(line_mesh, quad, comms_ctx, true) @test Spaces.column(horz_plane_space, 1, 1) isa Spaces.PointSpace # 2D Space - rectangle_mesh = CA.periodic_rectangle_mesh(; + rectangle_mesh = periodic_rectangle_mesh(; x_max = xlim[2], y_max = xlim[2], x_elem = helem, @@ -354,30 +354,26 @@ end @testset "make hybrid spaces" begin (; cent_space, face_space, xlim, zlim, velem, helem, npoly, quad) = get_cartesian_spaces() - config = CA.AtmosConfig( - Dict("topography" => "NoWarp", "topo_smoothing" => false), - ) device = ClimaComms.CPUSingleThreaded() - comms_ctx = ClimaComms.context(device) - z_stretch = Meshes.Uniform() - rectangle_mesh = CA.periodic_rectangle_mesh(; - x_max = xlim[2], - y_max = xlim[2], + context = ClimaComms.context(device) + grid = CA.BoxGrid( + Float32; + context, x_elem = helem, + x_max = xlim[2], y_elem = helem, + y_max = xlim[2], + z_elem = velem, + z_max = zlim[2], + nh_poly = npoly, + z_stretch = false, + bubble = true, + periodic_x = true, + periodic_y = true, ) - horz_plane_space = - CA.make_horizontal_space(rectangle_mesh, quad, comms_ctx, true) - test_cent_space, test_face_space = CA.make_hybrid_spaces( - horz_plane_space, - zlim[2], - velem, - z_stretch; - deep = false, - parsed_args = config.parsed_args, - ) - @test test_cent_space == cent_space - @test test_face_space == face_space + (; center_space, face_space) = CA.get_spaces(grid, context) + @test center_space == cent_space + @test face_space == face_space end @testset "promote_period" begin