diff --git a/docs/make.jl b/docs/make.jl index f289e9e7bf..2384b135c3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -25,6 +25,7 @@ makedocs(; "Home" => "index.md", "Installation instructions" => "installation_instructions.md", "Contributor Guide" => "contributor_guide.md", + "Performing Simulations" => "atmos_simulations.md", "Equations" => "equations.md", "EDMF Equations" => "edmf_equations.md", "Diagnostics" => "diagnostics.md", diff --git a/docs/src/atmos_simulations.md b/docs/src/atmos_simulations.md new file mode 100644 index 0000000000..b4a5d3c8a9 --- /dev/null +++ b/docs/src/atmos_simulations.md @@ -0,0 +1,38 @@ +# Performing a `ClimaAtmos` simulation + +`ClimaAtmos` supports a script-based interface. With this interface, simulations +are defined in a Julia script which uses `ClimaAtmos` as a library. Simulation +scripts can also contain post-processing and visualization. + +## The computational grid + +One of the first step in performing a simulation is defining the computational +grid. `ClimaAtmos` comes with a variety of pre-defined grids that are commonly +used in the field: + +- `StretchedColumnGrid`, a single column (1D) domain with non-uniform resolution + (higher resolution at the bottom of the column). The variant with constant + resolution is called `UniformColumnGrid`. Columns can only be run on + single-threaded simulations. + +By convention, all the `AbstractAtmosGrid`s in `ClimaAtmos` have a name that +ends in `Grid`. + +`AbstractAtmosGrid`s support a variety of features. For instance, to see an overview of +the computational grid: +```julia +println(summary(column)) # => +# Grid: UniformColumnGrid +# Number of elements: 10 +# Height: 30000.0 meters +# Grid stretching: Uniform +``` + +Users interested in adding new grids can do so by defining a new concrete +subtype `MyGrid` of the abstract `AbstractAtmosGrid` type. The only requirement +for `MyGrid` is that it has to have at least two fields: `center_space` and +`face_space`, which are `ClimaCore.Spaces` for the center and the face of the +cells respectively. We refer users to the `ClimaCore` documentation to learn +more about the notion of `Spaces`. + +Developers are encouraged to also define a `Base.summary` method. diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 752198d1b2..104772d5d4 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -121,4 +121,6 @@ include(joinpath("solver", "solve.jl")) include(joinpath("parameters", "create_parameters.jl")) +include(joinpath("simulation", "AtmosGrids.jl")) + end # module diff --git a/src/simulation/AtmosGrids.jl b/src/simulation/AtmosGrids.jl new file mode 100644 index 0000000000..11cd9147a0 --- /dev/null +++ b/src/simulation/AtmosGrids.jl @@ -0,0 +1,43 @@ +# AbstractAtmosGrids.jl +# +# In this file we define an abstract type AbstractAtmosGrid which describes the +# computational domain over which the simulation is run. We provide definitions of concrete +# AbstractAtmosGrids commonly used in simulations. + +module Grid + +import ClimaComms +import ClimaCore: Domains, Geometry, Meshes, Spaces, Topologies + +""" + AbstractAtmosGrid + + +Computational domain over which the simulation is run. Thin wrapper around two +`ClimaCore.Space`s, one the cell centers and one for the cell faces, and possibly other +general information and parameters related to the given grid. + +""" +abstract type AbstractAtmosGrid end + +""" + float_type(grid) + + +Return the floating-point type associated with the given `grid`. + +""" +float_type(grid::AbstractAtmosGrid) = Spaces.undertype(grid.center_space) + +""" + context(grid) + + +Return the context of the computational device associated with the given `grid`. + +""" +context(grid::AbstractAtmosGrid) = ClimaComms.context(grid.center_space) + +include("atmos_grids.jl") + +end diff --git a/src/simulation/atmos_grids.jl b/src/simulation/atmos_grids.jl new file mode 100644 index 0000000000..13ef308a98 --- /dev/null +++ b/src/simulation/atmos_grids.jl @@ -0,0 +1,139 @@ +# This file contains the definitions of common AbstractAtmosGrids. +# - ColumnGrid (with UniformColumnGrid and StretchedColumnGrid) +# +# We also provide convenience functions to build these grids. +include("atmos_grids_makers.jl") + +Base.@kwdef struct ColumnGrid{ + CS <: Spaces.ExtrudedFiniteDifferenceSpace, + FS <: Spaces.ExtrudedFiniteDifferenceSpace, + I <: Integer, + FT <: Real, + SR <: Meshes.StretchingRule, +} <: AbstractAtmosGrid + center_space::CS + face_space::FS + + z_elem::I + z_max::FT + + z_stretch::SR +end + +function Base.summary(io::IO, grid::ColumnGrid) + println(io, "Grid type: $(nameof(typeof(grid)))") + println(io, "Number of elements: $(grid.z_elem)") + println(io, "Height: $(grid.z_max) meters") + println(io, "Grid stretching: $(nameof(typeof(grid.z_stretch)))") + # Add information about the stretching, if any + for field in fieldnames(typeof(grid.z_stretch)) + println(io, " with: $(field): $(getproperty(grid.z_stretch, field))") + end +end + + +""" +function UniformColumnGrid(; z_elem, + z_max, + comms_ctx = ClimaComms.context(), + float_type = Float64) + +Construct an `ColumnGrid` for a column with uniform resolution. + +Keyword arguments +================= + +- `z_elem`: Number of elements. +- `z_max`: Height of the column (in meters). +- `comms_ctx`: Context of the computational environment where the simulation should be run, + as defined in ClimaComms. By default, the CLIMACOMMS_DEVICE environment + variable is read for one of "CPU", "CPUSingleThreaded", "CPUMultiThreaded", + "CUDA". If none is found, the fallback is to use a GPU (if available), or a + single threaded CPU (if not). +- `float_type`: Floating point type. Typically, Float32 or Float64 (default). + +""" +function UniformColumnGrid(; + z_elem, + z_max, + comms_ctx = ClimaComms.context(), + float_type = Float64, +) + isa(comms_ctx, ClimaComms.SingletonCommsContext) || + error("ColumnGrids are incompatible with MPI") + + # Promote types + z_max = float_type(z_max) + + # Vertical space + z_stretch = Meshes.Uniform() + z_space = + make_vertical_space(; z_elem, z_max, z_stretch, comms_ctx, float_type) + + # Horizontal space + h_space = make_trivial_horizontal_space(; comms_ctx, float_type) + + # 3D space + center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space) + face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space) + + return ColumnGrid(; center_space, face_space, z_elem, z_max, z_stretch) +end + + +""" +function StretchedColumnGrid(; z_elem, + dz_bottom, + dz_top, + z_max, + comms_ctx = ClimaComms.context(), + float_type = Float64) + +Construct an `AbstractAtmosGrid` for a column with non-uniform resolution (as +prescribed by `ClimaCore.Meshes.GeneralizedExponentialStretching`). + +Keyword arguments +================= + +- `z_elem`: Number of elements. +- `dz_bottom`: Resolution at the lower end of the column (in meters). +- `dz_top`: Resolution at the top end of the column (in meters). +- `z_max`: Height of the column (in meters). +- `comms_ctx`: Context of the computational environment where the simulation should be run, + as defined in ClimaComms. By default, the CLIMACOMMS_DEVICE environment + variable is read for one of "CPU", "CPUSingleThreaded", "CPUMultiThreaded", + "CUDA". If none is found, the fallback is to use a GPU (if available), or a + single threaded CPU (if not). +- `float_type`: Floating point type. Typically, Float32 or Float64 (default). + +""" +function StretchedColumnGrid(; + z_elem, + z_max, + dz_bottom, + dz_top, + comms_ctx = ClimaComms.context(), + float_type = Float64, +) + + isa(comms_ctx, ClimaComms.SingletonCommsContext) || + error("ColumnGrids are incompatible with MPI") + + # Promote types + z_max, dz_bottom, dz_top = + [float_type(v) for v in [z_max, dz_bottom, dz_top]] + + # Vertical space + z_stretch = Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top) + z_space = + make_vertical_space(; z_elem, z_max, z_stretch, comms_ctx, float_type) + + # Horizontal space + h_space = make_trivial_horizontal_space(; comms_ctx, float_type) + + # 3D space + center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space) + face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space) + + return ColumnGrid(; center_space, face_space, z_max, z_elem, z_stretch) +end diff --git a/src/simulation/atmos_grids_makers.jl b/src/simulation/atmos_grids_makers.jl new file mode 100644 index 0000000000..2c5592f3d8 --- /dev/null +++ b/src/simulation/atmos_grids_makers.jl @@ -0,0 +1,67 @@ +""" + make_vertical_space(; + z_elem, + z_max, + z_stretch::Meshes.StretchingRule, + comms_ctx::ClimaComms.AbstractCommsContext, + float_type + ) + +Return a vertical `Spaces.CenterFiniteDifferenceSpace` with `z_elem` and height `z_max` (in +meters) with resolution defined by `z_stretch`. +""" +function make_vertical_space(; + z_elem, + z_max, + z_stretch::Meshes.StretchingRule, + comms_ctx::ClimaComms.AbstractCommsContext, + float_type, +) + + # Promote types + z_max = float_type(z_max) + + z_domain = Domains.IntervalDomain( + Geometry.ZPoint(zero(float_type)), + Geometry.ZPoint(z_max); + boundary_tags = (:bottom, :top), + ) + z_mesh = Meshes.IntervalMesh(z_domain, z_stretch; nelems = z_elem) + z_topology = Topologies.IntervalTopology(comms_ctx, z_mesh) + return Spaces.CenterFiniteDifferenceSpace(z_topology) +end + + +""" + make_trivial_horizontal_space(; + comms_ctx::ClimaComms.AbstractCommsContext, + float_type + ) + +Return a one-point wide 2D `Spaces.SpectralElementSpace2D`. This is useful for building columns. +""" +function make_trivial_horizontal_space(; comms_ctx, float_type) + FT = float_type + + # 1-point wide horizontal domain + x_elem = y_elem = 1 + x_domain = Domains.IntervalDomain( + Geometry.XPoint(zero(FT)), + Geometry.XPoint(one(FT)); + periodic = true, + ) + y_domain = Domains.IntervalDomain( + Geometry.YPoint(zero(FT)), + Geometry.YPoint(one(FT)); + periodic = true, + ) + h_domain = Domains.RectangleDomain(x_domain, y_domain) + h_mesh = Meshes.RectilinearMesh(h_domain, x_elem, y_elem) + h_quadrature = Spaces.Quadratures.GL{1}() + h_topology = Topologies.Topology2D( + comms_ctx, + h_mesh, + Topologies.spacefillingcurve(h_mesh), + ) + return Spaces.SpectralElementSpace2D(h_topology, h_quadrature;) +end diff --git a/test/get_spaces.jl b/test/get_spaces.jl new file mode 100644 index 0000000000..fa4419e405 --- /dev/null +++ b/test/get_spaces.jl @@ -0,0 +1,266 @@ +using ClimaCore: + Geometry, Domains, Meshes, Topologies, Spaces, Hypsography, Fields +using ClimaComms + +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) + return Meshes.EquiangularCubedSphere(domain, h_elem) +end + +function make_horizontal_space( + mesh, + quad, + comms_ctx::ClimaComms.SingletonCommsContext, + bubble, +) + if mesh isa Meshes.AbstractMesh1D + topology = Topologies.IntervalTopology(comms_ctx, mesh) + space = Spaces.SpectralElementSpace1D(topology, quad) + elseif mesh isa Meshes.AbstractMesh2D + topology = Topologies.Topology2D( + comms_ctx, + mesh, + Topologies.spacefillingcurve(mesh), + ) + space = Spaces.SpectralElementSpace2D( + topology, + quad; + enable_bubble = bubble, + ) + end + return space +end + +function make_horizontal_space(mesh, quad, comms_ctx, bubble) + if mesh isa Meshes.AbstractMesh1D + error("Distributed mode does not work with 1D horizontal spaces.") + elseif mesh isa Meshes.AbstractMesh2D + topology = Topologies.DistributedTopology2D( + comms_ctx, + mesh, + Topologies.spacefillingcurve(mesh), + ) + space = Spaces.SpectralElementSpace2D( + topology, + quad; + enable_bubble = bubble, + ) + end + return space +end + +function make_hybrid_spaces( + h_space, + z_max, + z_elem, + z_stretch; + surface_warp = nothing, + topo_smoothing = false, +) + z_domain = Domains.IntervalDomain( + Geometry.ZPoint(zero(z_max)), + Geometry.ZPoint(z_max); + boundary_tags = (:bottom, :top), + ) + z_mesh = Meshes.IntervalMesh(z_domain, z_stretch; nelems = z_elem) + @info "z heights" z_mesh.faces + if surface_warp == nothing + device = ClimaComms.device(h_space) + comms_ctx = ClimaComms.SingletonCommsContext(device) + z_topology = Topologies.IntervalTopology(comms_ctx, z_mesh) + z_space = Spaces.CenterFiniteDifferenceSpace(z_topology) + center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space) + face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space) + else + z_surface = surface_warp(Fields.coordinate_field(h_space)) + topo_smoothing ? Hypsography.diffuse_surface_elevation!(z_surface) : + nothing + z_face_space = Spaces.FaceFiniteDifferenceSpace(z_mesh) + face_space = Spaces.ExtrudedFiniteDifferenceSpace( + h_space, + z_face_space, + Hypsography.LinearAdaption(z_surface), + ) + center_space = Spaces.CenterExtrudedFiniteDifferenceSpace(face_space) + end + return center_space, face_space +end + +function get_spaces(parsed_args, planet_radius, comms_ctx) + + FT = typeof(planet_radius) + z_elem = Int(parsed_args["z_elem"]) + z_max = FT(parsed_args["z_max"]) + dz_bottom = FT(parsed_args["dz_bottom"]) + dz_top = FT(parsed_args["dz_top"]) + topography = parsed_args["topography"] + bubble = parsed_args["bubble"] + + @assert topography in ("NoWarp", "DCMIP200", "Earth", "Agnesi", "Schar") + if topography == "DCMIP200" + warp_function = topography_dcmip200 + elseif topography == "Agnesi" + warp_function = topography_agnesi + elseif topography == "Schar" + warp_function = topography_schar + elseif topography == "NoWarp" + warp_function = nothing + elseif topography == "Earth" + data_path = joinpath(topo_elev_dataset_path(), "ETOPO1_coarse.nc") + earth_spline = NCDataset(data_path) do data + zlevels = data["elevation"][:] + lon = data["longitude"][:] + lat = data["latitude"][:] + # Apply Smoothing + smooth_degree = Int(parsed_args["smoothing_order"]) + esmth = imfilter(zlevels, Kernel.gaussian(smooth_degree)) + linear_interpolation( + (lon, lat), + esmth, + extrapolation_bc = (Periodic(), Flat()), + ) + end + @info "Generated interpolation stencil" + warp_function = generate_topography_warp(earth_spline) + end + @info "Topography" topography + + + h_elem = parsed_args["h_elem"] + radius = planet_radius + center_space, face_space = if parsed_args["config"] == "sphere" + nh_poly = parsed_args["nh_poly"] + quad = Spaces.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.GeneralizedExponentialStretching(dz_bottom, dz_top) + else + Meshes.Uniform() + end + if warp_function == nothing + make_hybrid_spaces(h_space, z_max, z_elem, z_stretch) + else + make_hybrid_spaces( + h_space, + z_max, + z_elem, + z_stretch; + surface_warp = warp_function, + topo_smoothing = parsed_args["topo_smoothing"], + ) + end + elseif parsed_args["config"] == "column" # single column + @warn "perturb_initstate flag is ignored for single column configuration" + Δx = FT(1) # Note: This value shouldn't matter, since we only have 1 column. + quad = Spaces.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.GeneralizedExponentialStretching(dz_bottom, dz_top) + else + Meshes.Uniform() + end + make_hybrid_spaces(h_space, z_max, z_elem, z_stretch) + elseif parsed_args["config"] == "box" + nh_poly = parsed_args["nh_poly"] + quad = Spaces.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.GeneralizedExponentialStretching(dz_bottom, dz_top) + else + Meshes.Uniform() + end + make_hybrid_spaces( + h_space, + z_max, + z_elem, + z_stretch; + surface_warp = warp_function, + ) + elseif parsed_args["config"] == "plane" + nh_poly = parsed_args["nh_poly"] + quad = Spaces.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.GeneralizedExponentialStretching(dz_bottom, dz_top) + else + Meshes.Uniform() + end + make_hybrid_spaces( + h_space, + z_max, + z_elem, + z_stretch; + surface_warp = warp_function, + ) + end + ncols = Fields.ncolumns(center_space) + ndofs_total = ncols * z_elem + hspace = Spaces.horizontal_space(center_space) + quad_style = Spaces.quadrature_style(hspace) + Nq = Spaces.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 diff --git a/test/runtests.jl b/test/runtests.jl index 72b455285e..8c1b3081e5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,6 +3,7 @@ using Test #! format: off @safetestset "Aqua" begin @time include("aqua.jl") end +@safetestset "AtmosGrids" begin @time include("atmos_grids.jl") end @safetestset "Utilities" begin @time include("utilities.jl") end @safetestset "Parameter tests" begin @time include("parameter_tests.jl") end @safetestset "Coupler Compatibility" begin @time include("coupler_compatibility.jl") end