Merge pull request #2103 from CliMA/gb/staggering_ergonomics

Add `face/center_space`

Author: charleskawczynski

NEWS.md

@@ -18,6 +18,13 @@ main
 - Fixed writing/reading purely vertical spaces. PR [2102](https://github.com/CliMA/ClimaCore.jl/pull/2102)
 - Fixed correctness bug in reductions on GPUs. PR [2106](https://github.com/CliMA/ClimaCore.jl/pull/2106)
 
+### ![][badge-✨feature/enhancement] `face_space`, `center_space` functions
+
+`ClimaCore.Spaces` now comes with two functions, `face_space` and
+`center_space`, to convert a `Space` from being cell-centered to be
+face-centered (and viceversa). These functions only work for vertical and
+extruded spaces.
+
 v0.14.20
 --------
 

src/CommonSpaces/CommonSpaces.jl

@@ -8,15 +8,20 @@ argument, `staggering::Staggering` to construct the desired space.
 module CommonSpaces
 
 export ExtrudedCubedSphereSpace,
-    CubedSphereSpace, ColumnSpace, Box3DSpace, SliceXZSpace, RectangleXYSpace
+    CubedSphereSpace,
+    ColumnSpace,
+    Box3DSpace,
+    SliceXZSpace,
+    RectangleXYSpace,
+    CellCenter,
+    CellFace
 
-export Grids
 import ClimaComms
 
 import ..DataLayouts,
     ..Meshes, ..Topologies, ..Geometry, ..Domains, ..Quadratures, ..Grids
 
-import ..Grids: Staggering
+import ..Grids: Staggering, CellCenter, CellFace
 import ..Spaces
 import ..CommonGrids
 import ..CommonGrids:
@@ -78,6 +83,19 @@ Note that these arguments are all the same as
 
 # Example usage
 
+```julia
+using ClimaCore.CommonSpaces
+space = ExtrudedCubedSphereSpace(;
+    z_elem = 10,
+    z_min = 0,
+    z_max = 1,
+    radius = 10,
+    h_elem = 10,
+    n_quad_points = 4
+    staggering = CellCenter()
+)
+```
+This will construct a cell-center space. If you wish to create a face centered space:
 ```julia
 using ClimaCore.CommonSpaces
 space = ExtrudedCubedSphereSpace(;
@@ -87,9 +105,10 @@ space = ExtrudedCubedSphereSpace(;
     radius = 10,
     h_elem = 10,
     n_quad_points = 4,
-    staggering = Grids.CellCenter()
+    staggering = CellFace()
 )
 ```
+alternatively, you can use the `Spaces.face_space` function.
 """
 function ExtrudedCubedSphereSpace end
 
@@ -186,7 +205,7 @@ space = ColumnSpace(;
     z_elem = 10,
     z_min = 0,
     z_max = 10,
-    staggering = Grids.CellCenter()
+    staggering = CellCenter()
 )
 ```
 """
@@ -270,7 +289,7 @@ space = Box3DSpace(;
     n_quad_points = 4,
     x_elem = 3,
     y_elem = 4,
-    staggering = Grids.CellCenter()
+    staggering = CellCenter()
 )
 ```
 """
@@ -319,8 +338,7 @@ configuration, given:
  - `quad` the quadrature style (defaults to `Quadratures.GLL{n_quad_points}`)
  - `staggering` vertical staggering, can be one of [[`Grids.CellFace`](@ref), [`Grids.CellCenter`](@ref)]
 
-Note that these arguments are all the same
-as [`CommonGrids.SliceXZGrid`](@ref),
+Note that these arguments are all the same as [`CommonGrids.SliceXZGrid`](@ref),
 except for `staggering`.
 
 # Example usage
@@ -336,7 +354,7 @@ space = SliceXZSpace(;
     periodic_x = false,
     n_quad_points = 4,
     x_elem = 4,
-    staggering = Grids.CellCenter()
+    staggering = CellCenter()
 )
 ```
 """
@@ -382,9 +400,6 @@ configuration, given:
  - `global_geometry` the global geometry (defaults to [`Geometry.CartesianGlobalGeometry`](@ref))
  - `quad` the quadrature style (defaults to `Quadratures.GLL{n_quad_points}`)
 
-Note that these arguments are all the same as [`CommonGrids.RectangleXYGrid`]
-(@ref), except for `staggering`.
-
 # Example usage
 
 ```julia

src/Spaces/Spaces.jl

@@ -80,10 +80,6 @@ global_geometry(space::AbstractSpace) = global_geometry(grid(space))
 space(refspace::AbstractSpace, staggering::Staggering) =
     space(grid(refspace), staggering)
 
-
-
-
-
 issubspace(::AbstractSpace, ::AbstractSpace) = false
 
 undertype(space::AbstractSpace) =
@@ -103,6 +99,14 @@ include("triangulation.jl")
 include("dss.jl")
 
 
+function center_space(space::AbstractSpace)
+    error("`center_space` can only be called with vertical/extruded spaces")
+end
+
+function face_space(space::AbstractSpace)
+    error("`center_space` can only be called with vertical/extruded spaces")
+end
+
 weighted_jacobian(space::Spaces.AbstractSpace) = local_geometry_data(space).WJ
 
 """

src/Spaces/extruded.jl

@@ -33,6 +33,29 @@ const FaceExtrudedFiniteDifferenceSpace{G} =
 const CenterExtrudedFiniteDifferenceSpace{G} =
     ExtrudedFiniteDifferenceSpace{G, CellCenter}
 
+"""
+    face_space(space::ExtrudedFiniteDifferenceSpace)
+
+Return face-centered space corresponding to `space`.
+
+If `space` is already face-centered, return itself.
+"""
+function face_space(space::ExtrudedFiniteDifferenceSpace)
+    return ExtrudedFiniteDifferenceSpace(grid(space), CellFace())
+end
+
+"""
+    center_space(space::ExtrudedFiniteDifferenceSpace)
+
+Return center-centered space corresponding to `space`.
+
+If `space` is already center-centered, return itself.
+"""
+function center_space(space::ExtrudedFiniteDifferenceSpace)
+    return ExtrudedFiniteDifferenceSpace(grid(space), CellCenter())
+end
+
+
 #=
 ExtrudedFiniteDifferenceSpace{S}(
     grid::Grids.ExtrudedFiniteDifferenceGrid,

src/Spaces/finitedifference.jl

@@ -81,7 +81,27 @@ CenterFiniteDifferenceSpace(
 Adapt.adapt_structure(to, space::FiniteDifferenceSpace) =
     FiniteDifferenceSpace(Adapt.adapt(to, grid(space)), staggering(space))
 
+"""
+    face_space(space::FiniteDifferenceSpace)
+
+Return face-centered space corresponding to `space`.
+
+If `space` is already face-centered, return itself.
+"""
+function face_space(space::FiniteDifferenceSpace)
+    return FiniteDifferenceSpace(grid(space), CellFace())
+end
 
+"""
+    center_space(space::FiniteDifferenceSpace)
+
+Return center-centered space corresponding to `space`.
+
+If `space` is already center-centered, return itself.
+"""
+function center_space(space::FiniteDifferenceSpace)
+    return FiniteDifferenceSpace(grid(space), CellCenter())
+end
 
 nlevels(space::FiniteDifferenceSpace) = length(space)
 # TODO: deprecate?

test/Spaces/unit_spaces.jl

@@ -112,6 +112,12 @@ on_gpu || @testset "extruded (2d 1×3) finite difference space" begin
     # Extrusion
     f_space = Spaces.ExtrudedFiniteDifferenceSpace(hspace, vert_face_space)
     c_space = Spaces.CenterExtrudedFiniteDifferenceSpace(f_space)
+
+    @test f_space == Spaces.face_space(f_space)
+    @test c_space == Spaces.center_space(f_space)
+    @test f_space == Spaces.face_space(c_space)
+    @test c_space == Spaces.center_space(c_space)
+
     s = DataLayouts.farray_size(Spaces.coordinates_data(c_space))
     z = Fields.coordinate_field(c_space).z
     @test s == (10, 4, 2, 5) # 10V, 4I, 2F(x,z), 5H
@@ -145,19 +151,25 @@ end
     mesh = Meshes.IntervalMesh(domain; nelems = 1)
     topology = Topologies.IntervalTopology(context, mesh)
 
-    space = Spaces.CenterFiniteDifferenceSpace(topology)
-    @test repr(space) == """
+    c_space = Spaces.CenterFiniteDifferenceSpace(topology)
+    f_space = Spaces.FaceFiniteDifferenceSpace(topology)
+    @test repr(c_space) == """
     CenterFiniteDifferenceSpace:
       context: SingletonCommsContext using CPUSingleThreaded
       mesh: 1-element IntervalMesh of IntervalDomain: z ∈ [0.0,5.0] (:bottom, :top)"""
 
-    coord_data = Spaces.coordinates_data(space)
-    point_space = Spaces.level(space, 1)
+    @test f_space == Spaces.face_space(f_space)
+    @test c_space == Spaces.center_space(f_space)
+    @test f_space == Spaces.face_space(c_space)
+    @test c_space == Spaces.center_space(c_space)
+
+    coord_data = Spaces.coordinates_data(c_space)
+    point_space = Spaces.level(c_space, 1)
     @test point_space isa Spaces.PointSpace
     @test Spaces.coordinates_data(point_space)[] ==
           Spaces.level(coord_data, 1)[]
 
-    @test Spaces.local_geometry_type(typeof(space)) <: Geometry.LocalGeometry
+    @test Spaces.local_geometry_type(typeof(c_space)) <: Geometry.LocalGeometry
 
     x_max = FT(1)
     y_max = FT(1)
@@ -186,7 +198,7 @@ end
     @test length(Spaces.all_nodes(hspace)) == 4
 
     if on_gpu
-        adapted_space = adapt(space)(space)
+        adapted_space = adapt(c_space)(c_space)
         @test ClimaComms.context(adapted_space) == DeviceSideContext()
         @test ClimaComms.device(adapted_space) == DeviceSideDevice()