Mv ncolumns to Spaces, define nlevels on fields (#2199)

* Mv ncolumns to Spaces, define nlevels on fields

* Update src/Spaces/Spaces.jl

Co-authored-by: Gabriele Bozzola <[email protected]>

---------

Co-authored-by: Gabriele Bozzola <[email protected]>

Author: charleskawczynski

src/Fields/Fields.jl

@@ -21,6 +21,7 @@ import ..Topologies
 import ..Quadratures
 import ..Grids: ColumnIndex, local_geometry_type
 import ..Spaces: Spaces, AbstractSpace, AbstractPointSpace, cuda_synchronize
+import ..Spaces: nlevels, ncolumns
 import ..Geometry: Geometry, Cartesian12Vector
 import ..Utilities: PlusHalf, half, UnrolledFunctions
 

src/Fields/indices.jl

@@ -107,25 +107,17 @@ bycolumn(
 
 """
     ncolumns(::Field)
-    ncolumns(::Space)
 
-Number of columns in a given space.
+Number of columns in the field's space.
 """
 ncolumns(field::Field) = ncolumns(axes(field))
 
-ncolumns(space::Spaces.ExtrudedFiniteDifferenceSpace) =
-    ncolumns(Spaces.horizontal_space(space))
+"""
+    nlevels(::Field)
 
-function ncolumns(space::Spaces.SpectralElementSpace1D)
-    Nh = Topologies.nlocalelems(space)
-    Nq = Quadratures.degrees_of_freedom(Spaces.quadrature_style(space))
-    return Nh * Nq
-end
-function ncolumns(space::Spaces.SpectralElementSpace2D)
-    Nh = Topologies.nlocalelems(space)
-    Nq = Quadratures.degrees_of_freedom(Spaces.quadrature_style(space))
-    return Nh * Nq * Nq
-end
+Number of levels in the field's space.
+"""
+nlevels(field::Field) = nlevels(axes(field))
 
 # potential TODO:
 # - define a ColumnIndices type, make it work with https://github.com/JuliaFolds/FLoops.jl

src/Spaces/Spaces.jl

@@ -141,8 +141,8 @@ ClimaComms.array_type(space::AbstractSpace) =
 
 The domain maximum along the z-direction.
 """
-function z_max(space::Spaces.AbstractSpace)
-    mesh = Topologies.mesh(Spaces.vertical_topology(space))
+function z_max(space::AbstractSpace)
+    mesh = Topologies.mesh(vertical_topology(space))
     domain = Topologies.domain(mesh)
     return Domains.z_max(domain)
 end
@@ -152,10 +152,30 @@ end
 
 The domain minimum along the z-direction.
 """
-function z_min(space::Spaces.AbstractSpace)
-    mesh = Topologies.mesh(Spaces.vertical_topology(space))
+function z_min(space::AbstractSpace)
+    mesh = Topologies.mesh(vertical_topology(space))
     domain = Topologies.domain(mesh)
     return Domains.z_min(domain)
 end
 
+"""
+    ncolumns(space::AbstractSpace)
+
+Number of columns in a given `space` on the local processor.
+"""
+ncolumns(space::ExtrudedFiniteDifferenceSpace) =
+    ncolumns(horizontal_space(space))
+
+function ncolumns(space::SpectralElementSpace1D)
+    Nh = Topologies.nlocalelems(space)
+    Nq = Quadratures.degrees_of_freedom(quadrature_style(space))
+    return Nh * Nq
+end
+function ncolumns(space::SpectralElementSpace2D)
+    Nh = Topologies.nlocalelems(space)
+    Nq = Quadratures.degrees_of_freedom(quadrature_style(space))
+    return Nh * Nq * Nq
+end
+
+
 end # module