From 447b36be78a5d944a734c6cc3f6547745e1c7638 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Tue, 3 Oct 2023 10:51:33 -0700 Subject: [PATCH] Add SphereGrid [skip ci][ci skip] --- docs/src/atmos_simulations.md | 8 +- src/simulation/atmos_grids.jl | 205 ++++++++++++++++++- test/atmos_grids.jl | 371 ++++++++++++++++++++++++++++++++++ 3 files changed, 582 insertions(+), 2 deletions(-) create mode 100644 test/atmos_grids.jl diff --git a/docs/src/atmos_simulations.md b/docs/src/atmos_simulations.md index f46812aa426..5ead091ba3f 100644 --- a/docs/src/atmos_simulations.md +++ b/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`. diff --git a/src/simulation/atmos_grids.jl b/src/simulation/atmos_grids.jl index ad3621d6f62..e640e917ec4 100644 --- a/src/simulation/atmos_grids.jl +++ b/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 diff --git a/test/atmos_grids.jl b/test/atmos_grids.jl new file mode 100644 index 00000000000..b695cc3877a --- /dev/null +++ b/test/atmos_grids.jl @@ -0,0 +1,371 @@ +using Test +import ClimaComms +import ClimaCore: Meshes +import ClimaAtmos as CA + +# Here, we copy and paste the ClimaAtmos.get_spaces functions from commit fab2f3c to ensure +# that we produce the same spaces. (The function is modified to take planet_radius directly, +# otherwise it would be a massive pain) +# include("get_spaces.jl") + +@testset "UniformColumnGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32) + z_elem = 10 + z_max = 30000 + + column = CA.Grid.UniformColumnGrid(; z_elem, z_max, float_type = FT) + @test column.z_max == convert(FT, z_max) + @test column.z_stretch == Meshes.Uniform() + @test CA.Grid.float_type(column) == FT + @test CA.Grid.context(column) == ClimaComms.context() + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "z_elem" => z_elem, + "z_max" => z_max, + "config" => "column", + # Not used + "h_elem" => NaN, + "z_stretch" => false, + "dz_bottom" => NaN, + "dz_top" => NaN, + "topography" => "NoWarp", + "bubble" => false, + ) + planet_radius = convert(FT, 1.0) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(column.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(column.face_space) == typeof(traditional_spaces.face_space) + + @test parent(column.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(column.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + + end +end + +@testset "StretchedColumnGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32) + dz_bottom = 500 + dz_top = 5000 + z_elem = 10 + z_max = 30000 + + column = CA.Grid.StretchedColumnGrid(; + z_elem, + z_max, + dz_bottom, + dz_top, + float_type = FT, + ) + @test column.z_max == convert(FT, z_max) + @test column.z_stretch == + Meshes.GeneralizedExponentialStretching(FT(dz_bottom), FT(dz_top)) + @test CA.Grid.float_type(column) == FT + @test CA.Grid.context(column) == ClimaComms.context() + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "z_elem" => z_elem, + "z_max" => z_max, + "dz_bottom" => dz_bottom, + "dz_top" => dz_top, + "z_stretch" => true, + "config" => "column", + # Not used + "h_elem" => NaN, + "topography" => "NoWarp", + "bubble" => false, + ) + planet_radius = convert(FT, 1.0) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(column.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(column.face_space) == typeof(traditional_spaces.face_space) + + @test parent(column.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(column.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + end +end + +@testset "VerticallyStretchedBoxGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32), enable_bubble in (true, false) + nh_poly = 3 + dz_bottom = 500 + dz_top = 5000 + x_elem = 6 + y_elem = 6 + x_max = 300000 + y_max = 300000 + z_elem = 10 + z_max = 30000 + + box = CA.Grid.Box(; + nh_poly, + x_elem, + y_elem, + z_elem, + x_max, + y_max, + z_max, + dz_bottom, + dz_top, + float_type = FT, + enable_bubble, + ) + @test box.x_max == convert(FT, x_max) + @test box.y_max == convert(FT, y_max) + @test box.z_max == convert(FT, z_max) + @test CA.Grid.float_type(box) == FT + @test CA.Grid.context(box) == ClimaComms.context() + @test box.z_stretch == + Meshes.GeneralizedExponentialStretching(FT(dz_bottom), FT(dz_top)) + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "nh_poly" => nh_poly, + "x_elem" => x_elem, + "x_max" => x_max, + "y_elem" => y_elem, + "y_max" => y_max, + "z_elem" => z_elem, + "z_max" => z_max, + "dz_bottom" => dz_bottom, + "dz_top" => dz_top, + "z_stretch" => true, + "config" => "box", + "bubble" => enable_bubble, + # Not used + "h_elem" => NaN, + "topography" => "NoWarp", + ) + planet_radius = convert(FT, 1.0) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(box.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(box.face_space) == typeof(traditional_spaces.face_space) + + @test parent(box.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(box.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + end +end + +@testset "VerticallyUniformBoxGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32), enable_bubble in (true, false) + nh_poly = 3 + x_elem = 6 + y_elem = 6 + x_max = 300000 + y_max = 300000 + z_elem = 10 + z_max = 30000 + + box = CA.Grid.VerticallyUniformBoxGrid(; + nh_poly, + x_elem, + y_elem, + z_elem, + x_max, + y_max, + z_max, + float_type = FT, + enable_bubble, + ) + @test box.x_max == convert(FT, x_max) + @test box.y_max == convert(FT, y_max) + @test box.z_max == convert(FT, z_max) + @test CA.Grid.float_type(box) == FT + @test CA.Grid.context(box) == ClimaComms.context() + @test box.z_stretch == Meshes.Uniform() + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "nh_poly" => nh_poly, + "x_elem" => x_elem, + "x_max" => x_max, + "y_elem" => y_elem, + "y_max" => y_max, + "z_elem" => z_elem, + "z_max" => z_max, + "z_stretch" => false, + "config" => "box", + "bubble" => enable_bubble, + # Not used + "dz_bottom" => NaN, + "dz_top" => NaN, + "h_elem" => NaN, + "topography" => "NoWarp", + ) + planet_radius = convert(FT, 1.0) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(box.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(box.face_space) == typeof(traditional_spaces.face_space) + + @test parent(box.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(box.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + end +end + +@testset "VerticallyStretchedSphereGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32), enable_bubble in (true, false) + nh_poly = 3 + dz_bottom = 500 + dz_top = 5000 + h_elem = 6 + radius = 100 + z_elem = 10 + z_max = 30000 + + sphere = CA.Grid.Sphere(; + nh_poly, + h_elem, + z_elem, + radius, + z_max, + dz_bottom, + dz_top, + float_type = FT, + enable_bubble, + ) + @test sphere.radius == convert(FT, radius) + @test CA.Grid.float_type(sphere) == FT + @test CA.Grid.context(sphere) == ClimaComms.context() + @test sphere.z_stretch == + Meshes.GeneralizedExponentialStretching(FT(dz_bottom), FT(dz_top)) + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "nh_poly" => nh_poly, + "radius" => radius, + "h_elem" => h_elem, + "z_elem" => z_elem, + "z_max" => z_max, + "dz_bottom" => dz_bottom, + "dz_top" => dz_top, + "z_stretch" => true, + "config" => "sphere", + "bubble" => enable_bubble, + # Not used + "topography" => "NoWarp", + ) + planet_radius = convert(FT, radius) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(sphere.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(sphere.face_space) == typeof(traditional_spaces.face_space) + + @test parent(sphere.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(sphere.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + end +end + +@testset "VerticallyUniformSphereGrid" begin + + # Check that we can build some grids + for FT in (Float64, Float32), enable_bubble in (true, false) + nh_poly = 3 + h_elem = 6 + y_elem = 6 + x_max = 300000 + y_max = 300000 + z_elem = 10 + z_max = 30000 + + sphere = CA.Grid.VerticallyUniformSphereGrid(; + nh_poly, + h_elem, + radius, + z_elem, + z_max, + float_type = FT, + enable_bubble, + ) + @test sphere.radius == convert(FT, radius) + @test CA.Grid.float_type(sphere) == FT + @test CA.Grid.context(sphere) == ClimaComms.context() + @test sphere.z_stretch == Meshes.Uniform() + + # The traditional interface takes arguments even if they are not needed + parsed_args = Dict( + "nh_poly" => nh_poly, + "radius" => radius, + "h_elem" => h_elem, + "z_elem" => z_elem, + "z_max" => z_max, + "z_stretch" => false, + "config" => "sphere", + "bubble" => enable_bubble, + # Not used + "dz_bottom" => NaN, + "dz_top" => NaN, + "topography" => "NoWarp", + ) + planet_radius = convert(FT, radius) + + traditional_spaces = + get_spaces(parsed_args, planet_radius, ClimaComms.context()) + + # We cannot compare directly the spaces because they are complex objects made with + # mutable components, so we compare the relevant bits + + @test typeof(sphere.center_space) == + typeof(traditional_spaces.center_space) + @test typeof(sphere.face_space) == typeof(traditional_spaces.face_space) + + @test parent(sphere.center_space.center_local_geometry) == + parent(traditional_spaces.center_space.center_local_geometry) + @test parent(sphere.face_space.center_local_geometry) == + parent(traditional_spaces.face_space.center_local_geometry) + end +end