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 add a default staggering to `CommonSpaces` and a function to go from center to face spaces and back.
CliMA · Dec 16, 2024 · 5f44644 · 5f44644
1 parent 340603b
commit 5f44644
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 23 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
@@ -10,13 +10,12 @@ module CommonSpaces
 export ExtrudedCubedSphereSpace,
     CubedSphereSpace, ColumnSpace, Box3DSpace, SliceXZSpace, RectangleXYSpace
 
-export Grids
 import ClimaComms
 
 import ..DataLayouts,
     ..Meshes, ..Topologies, ..Geometry, ..Domains, ..Quadratures, ..Grids
 
-import ..Grids: Staggering
+import ..Grids: Staggering, CellCenter
 import ..Spaces
 import ..CommonGrids
 import ..CommonGrids:
@@ -80,24 +79,38 @@ Note that these arguments are all the same as
 
 ```julia
 using ClimaCore.CommonSpaces
+space = ExtrudedCubedSphereSpace(;
+    z_elem = 10,
+    z_min = 0,
+    z_max = 1,
+    radius = 10,
+    h_elem = 10,
+    n_quad_points = 4
+)
+```
+This will construct a cell-center space. If you wish to create a face centered space:
+```julia
+using ClimaCore.CommonSpaces
+import Grids: FaceCenter
 space = ExtrudedCubedSphereSpace(;
     z_elem = 10,
     z_min = 0,
     z_max = 1,
     radius = 10,
     h_elem = 10,
     n_quad_points = 4,
-    staggering = Grids.CellCenter()
+    staggering = FaceCenter()
 )
 ```
+alternatively, you can use the [`Spaces.face_space`](@ref) function.
 """
 function ExtrudedCubedSphereSpace end
 
 ExtrudedCubedSphereSpace(; kwargs...) =
     ExtrudedCubedSphereSpace(Float64; kwargs...)
 ExtrudedCubedSphereSpace(
     ::Type{FT};
-    staggering::Staggering,
+    staggering::Staggering = CellCenter(),
     kwargs...,
 ) where {FT} = Spaces.ExtrudedFiniteDifferenceSpace(
     ExtrudedCubedSphereGrid(FT; kwargs...),
@@ -185,14 +198,17 @@ using ClimaCore.CommonSpaces
 space = ColumnSpace(;
     z_elem = 10,
     z_min = 0,
-    z_max = 10,
-    staggering = Grids.CellCenter()
+    z_max = 10
 )
 ```
 """
 function ColumnSpace end
 ColumnSpace(; kwargs...) = ColumnSpace(Float64; kwargs...)
-ColumnSpace(::Type{FT}; staggering::Staggering, kwargs...) where {FT} =
+ColumnSpace(
+    ::Type{FT};
+    staggering::Staggering = Grids.CellCenter(),
+    kwargs...,
+) where {FT} =
     Spaces.FiniteDifferenceSpace(ColumnGrid(FT; kwargs...), staggering)
 
 """
@@ -269,14 +285,17 @@ space = Box3DSpace(;
     periodic_y = false,
     n_quad_points = 4,
     x_elem = 3,
-    y_elem = 4,
-    staggering = Grids.CellCenter()
+    y_elem = 4
 )
 ```
 """
 function Box3DSpace end
 Box3DSpace(; kwargs...) = Box3DSpace(Float64; kwargs...)
-Box3DSpace(::Type{FT}; staggering::Staggering, kwargs...) where {FT} =
+Box3DSpace(
+    ::Type{FT};
+    staggering::Staggering = Grids.CellCenter(),
+    kwargs...,
+) where {FT} =
     Spaces.ExtrudedFiniteDifferenceSpace(Box3DGrid(FT; kwargs...), staggering)
 
 """
@@ -335,14 +354,17 @@ space = SliceXZSpace(;
     z_max = 1,
     periodic_x = false,
     n_quad_points = 4,
-    x_elem = 4,
-    staggering = Grids.CellCenter()
+    x_elem = 4
 )
 ```
 """
 function SliceXZSpace end
 SliceXZSpace(; kwargs...) = SliceXZSpace(Float64; kwargs...)
-SliceXZSpace(::Type{FT}; staggering::Staggering, kwargs...) where {FT} =
+SliceXZSpace(
+    ::Type{FT};
+    staggering::Staggering = Grids.CellCenter(),
+    kwargs...,
+) where {FT} =
     Spaces.ExtrudedFiniteDifferenceSpace(SliceXZGrid(FT; kwargs...), staggering)
 
 """

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()