Add face/center_space

As I started using the `CommonSpaces`, I realized that specifying the staggering was a pain point. It required importing `Grids`, and I pretty much always wanted to build CenterSpaces anyways. In addition to this, I realized that going from center to faces and viceversa was not trivial. This required carrying around both the center and the face space when only one is really needed. In this commit, I reexport the staggering in `CommonSpaces` and a function to go from center to face spaces and back.
CliMA · Dec 18, 2024 · 1f67e1f · 1f67e1f
1 parent 340603b
commit 1f67e1f
Show file tree

Hide file tree

Showing 6 changed files with 104 additions and 22 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -12,7 +12,15 @@ main
    - [SliceXZSpace]()
    - [RectangleXYSpace]()
 
+### ![][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
+--------
 
  - We've added new convenience constructors for grids PR [1848](https://github.com/CliMA/ClimaCore.jl/pull/1848). Here are links to the new constructors:
    - [ExtrudedCubedSphereGrid](https://github.com/CliMA/ClimaCore.jl/blob/cbb193042fac3b4bef33251fbc0f232427bfe506/src/CommonGrids/CommonGrids.jl#L85-L144)

diff --git a/src/CommonSpaces/CommonSpaces.jl b/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,
+    FaceCenter
 
-export Grids
 import ClimaComms
 
 import ..DataLayouts,
     ..Meshes, ..Topologies, ..Geometry, ..Domains, ..Quadratures, ..Grids
 
-import ..Grids: Staggering
+import ..Grids: Staggering, CellCenter, FaceCenter
 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 = FaceCenter()
 )
 ```
+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

diff --git a/src/Spaces/Spaces.jl b/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
 
 """

diff --git a/src/Spaces/extruded.jl b/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,

diff --git a/src/Spaces/finitedifference.jl b/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?

diff --git a/test/Spaces/unit_spaces.jl b/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()