Add AtmosGrids and ColumnGrid

[skip ci][ci skip]

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",

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.

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

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

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

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