Merge pull request #2330 from CliMA/dy/missing_methods

Add missing methods for column spaces and FieldVectors

Author: dennisYatunin

src/Fields/Fields.jl

@@ -30,7 +30,7 @@ import ..Utilities: PlusHalf, half
 using ..RecursiveApply
 using ClimaComms
 import Adapt
-import UnrolledUtilities: unrolled_map
+import UnrolledUtilities: unrolled_map, unrolled_findfirst
 
 import StaticArrays, LinearAlgebra, Statistics, InteractiveUtils
 

src/Fields/fieldvector.jl

@@ -459,6 +459,17 @@ import ClimaComms
 ClimaComms.array_type(x::FieldVector) =
     promote_type(unrolled_map(ClimaComms.array_type, _values(x))...)
 
+ClimaComms.device(x::FieldVector) = ClimaComms.device(ClimaComms.context(x))
+function ClimaComms.context(x::FieldVector)
+    isempty(_values(x)) && error("Empty FieldVector has no device or context")
+    # We don't have promotion for devices or contexts, so we use the first value
+    # that isn't a PointField (a PointField's data can be stored on a different
+    # device from other Fields to avoid scalar indexing on GPUs). If there is no
+    # such value, fall back to using the first PointField.
+    index = unrolled_findfirst(Base.Fix1(!isa, PointField), _values(x))
+    return ClimaComms.context(_values(x)[isnothing(index) ? 1 : index])
+end
+
 function __rprint_diff(
     io::IO,
     x::T,

src/Grids/column.jl

@@ -36,6 +36,8 @@ struct ColumnGrid{
     colidx::C
 end
 
+Adapt.@adapt_structure ColumnGrid
+
 local_geometry_type(::Type{ColumnGrid{G, C}}) where {G, C} =
     local_geometry_type(G)
 

src/Spaces/finitedifference.jl

@@ -103,6 +103,7 @@ function center_space(space::FiniteDifferenceSpace)
     return FiniteDifferenceSpace(grid(space), CellCenter())
 end
 
+ncolumns(::FiniteDifferenceSpace) = 1
 nlevels(space::FiniteDifferenceSpace) = length(space)
 # TODO: deprecate?
 Base.length(space::FiniteDifferenceSpace) = length(coordinates_data(space))