Add SphereGrid

[skip ci][ci skip]

docs/src/atmos_simulations.md

@@ -14,11 +14,17 @@ used in the field:
   (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.
-- `VerticallyStretchedBoxGrid`, a period box with columns (1D) domain with
+- `VerticallyStretchedBoxGrid`, a periodic box with columns (1D) domain with
   non-uniform resolution in the vertical direction (higher resolution at the
   bottom of the column). The variant with vertically uniform resolution is
   `VerticallyUniformBoxGrid`. Given that `VerticallyStretchedBoxGrid` is a
   commonly used domain, we also provide an alias `Box` for it.
+- `VerticallyStretchedSphereGrid`, a equiangular cubed sphere with columns (1D)
+  domain with non-uniform resolution in the vertical direction (higher
+  resolution at the bottom of the column). The variant with vertically uniform
+  resolution is `VerticallyUniformSphereGrid`. Given that
+  `VerticallyStretchedSphereGrid` is a commonly used domain, we also provide an
+  alias `Sphere` for it.
 
 By convention, all the `AbstractAtmosGrid`s in `ClimaAtmos` have a name that
 ends in `Grid`.

src/simulation/atmos_grids.jl

@@ -1,12 +1,14 @@
 # This file contains the definitions of common AbstractAtmosGrids.
 # - ColumnGrid (with UniformColumnGrid and StretchedColumnGrid)
 # - BoxGrid (with VerticallyUniformBoxGrid and VerticallyStretchedBoxGrid)
+# - SphereGrid (with VerticallyUniformSphereGrid and VerticallyStretchedSphereGrid)
 #
 # We provide aliases for common grids:
 # - Box = VerticallyStretchedBoxGrid
+# - Sphere = VerticallyStretchedSphereGrid
 #
 # We also provide convenience functions to build these grids.
-include("atmos_grids_makers.jl")
+# include("atmos_grids_makers.jl")
 
 ##############
 # ColumnGrid #
@@ -382,3 +384,204 @@ function VerticallyUniformBoxGrid(;
         enable_bubble,
     )
 end
+
+##############
+# SphereGrid #
+##############
+
+Base.@kwdef struct SphereGrid{
+    CS <: Spaces.ExtrudedFiniteDifferenceSpace,
+    FS <: Spaces.ExtrudedFiniteDifferenceSpace,
+    I <: Integer,
+    FT <: Real,
+    SR <: Meshes.StretchingRule,
+} <: AbstractAtmosGrid
+    center_space::CS
+    face_space::FS
+
+    nh_poly::I
+
+    h_elem::I
+    radius::FT
+    z_elem::I
+    z_max::FT
+
+    z_stretch::SR
+
+    enable_bubble::Bool
+end
+
+function Base.summary(io::IO, grid::SphereGrid)
+    println(io, "Grid type: $(nameof(typeof(grid)))")
+    println(io, "Number of vertical elements: $(grid.z_elem)")
+    println(io, "Height: $(grid.z_max) meters")
+    println(io, "Vertical 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
+    println(io, "Radius: $radius meters")
+    println(io, "Horizontal elements per edge")
+    println(io, "  with: $(grid.nh_poly)-degree polynomials")
+    println(
+        io,
+        "  ",
+        grid.enable_bubble ? "with" : "without",
+        " bubble correction",
+    )
+end
+
+"""
+function VerticallyStretchedSphereGrid(; nh_poly,
+                                         h_elem,
+                                         radius,
+                                         z_elem,
+                                         z_max,
+                                         dz_bottom,
+                                         dz_top,
+                                         enable_bubble = false,
+                                         comms_ctx = ClimaComms.context(),
+                                         float_type = Float64)
+
+Construct an `SphereGrid` for a cubed sphere with columns with non-uniform
+resolution (as prescribed by
+`ClimaCore.Meshes.GeneralizedExponentialStretching`).
+
+Keyword arguments
+=================
+
+- `nh_poly`: Horizontal polynomial degree.
+- `radius`: Radius of the sphere (in meters).
+- `h_elem`: Number of spectral elements per edge.
+- `x_elem`: Number of spectral elements on the x direction.
+- `x_max`: Length of the box (in meters) (`x_min` is 0).
+- `y_elem`: Number of spectral elements on the y direction.
+- `y_max`: Depth of the box (in meters) (`y_min` is 0).
+- `z_elem`: Number of spectral elements on the vertical direction.
+- `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).
+- `enable_bubble`: Enables the `bubble correction` for more accurate element areas.
+- `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 VerticallyStretchedSphereGrid(;
+    nh_poly,
+    h_elem,
+    radius,
+    z_elem,
+    z_max,
+    dz_bottom,
+    dz_top,
+    enable_bubble = false,
+    comms_ctx = ClimaComms.context(),
+    float_type = Float64,
+)
+
+    # Promote types
+    radius, dz_bottom, dz_top, z_max =
+        [float_type(v) for v in [radius, dz_bottom, dz_top, z_max]]
+
+    # Vertical
+    z_stretch = Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top)
+    z_space =
+        make_vertical_space(; z_elem, z_max, z_stretch, comms_ctx, float_type)
+
+    # Horizontal
+    h_domain = Domains.SphereDomain(radius)
+    h_mesh = Meshes.EquiangularCubedSphere(h_domain, h_elem)
+    h_space = make_horizontal_space(; nh_poly, h_mesh, comms_ctx, enable_bubble)
+
+    center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space)
+    face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space)
+
+    return SphereGrid(;
+        center_space,
+        face_space,
+        nh_poly,
+        h_elem,
+        radius,
+        z_elem,
+        z_max,
+        z_stretch,
+        enable_bubble,
+    )
+end
+
+# Alias for a commonly used grid type
+const Sphere = VerticallyStretchedSphereGrid
+
+"""
+function VerticallyUniformSphereGrid(; nh_poly,
+                                         h_elem,
+                                         radius,
+                                         z_elem,
+                                         z_max,
+                                         enable_bubble = false,
+                                         comms_ctx = ClimaComms.context(),
+                                         float_type = Float64)
+
+
+Construct an `SphereGrid` for a cubed sphere with columns with uniform resolution.
+
+Keyword arguments
+=================
+
+- `nh_poly`: Horizontal polynomial degree.
+- `radius`: Radius of the sphere (in meters).
+- `h_elem`: Number of spectral elements per edge.
+- `z_elem`: Number of spectral elements on the vertical direction.
+- `z_max`: Height of the column (in meters).
+- `enable_bubble`: Enables the `bubble correction` for more accurate element areas.
+- `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 VerticallyUniformSphereGrid(;
+    nh_poly,
+    h_elem,
+    radius,
+    z_elem,
+    z_max,
+    enable_bubble = false,
+    comms_ctx = ClimaComms.context(),
+    float_type = Float64,
+)
+
+    # Promote types
+    radius, z_max = [float_type(v) for v in [radius, z_max]]
+
+    # Vertical
+    z_stretch = Meshes.Uniform()
+    z_space =
+        make_vertical_space(; z_elem, z_max, z_stretch, comms_ctx, float_type)
+
+    # Horizontal
+    h_domain = Domains.SphereDomain(radius)
+    h_mesh = Meshes.EquiangularCubedSphere(h_domain, h_elem)
+    h_space = make_horizontal_space(; nh_poly, h_mesh, comms_ctx, enable_bubble)
+
+    center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space)
+    face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space)
+
+    return SphereGrid(;
+        center_space,
+        face_space,
+        nh_poly,
+        h_elem,
+        radius,
+        z_elem,
+        z_max,
+        z_stretch,
+        enable_bubble,
+    )
+end