Specialize field constructor for boolean fields (#2239)

Author: charleskawczynski

NEWS.md

@@ -4,6 +4,8 @@ ClimaCore.jl Release Notes
 main
 -------
 
+ - `Fields.Field(Bool, ::AbstractSpace)` is now supported. PR [2239](https://github.com/CliMA/ClimaCore.jl/pull/2239).
+
  - `SpectralElementSpace2D` constructors now support nodal masks. PR [2201](https://github.com/CliMA/ClimaCore.jl/pull/2201). See its documentation [here](https://clima.github.io/ClimaCore.jl/dev/masks). Note that it does not yet support restarts.
 
 - Added support for InputOutput with PointSpaces

src/Fields/Fields.jl

@@ -53,6 +53,13 @@ Field(values::V, space::S) where {V <: AbstractData, S <: AbstractSpace} =
 Field(::Type{T}, space::S) where {T, S <: AbstractSpace} =
     Field(similar(Spaces.coordinates_data(space), T), space)
 
+function Field(::Type{Bool}, space::S) where {S <: AbstractSpace}
+    FT = Spaces.undertype(space)
+    data = similar(Spaces.coordinates_data(space), FT)
+    bool_data = DataLayouts.replace_basetype(data, Bool)
+    return Field(bool_data, space)
+end
+
 local_geometry_type(::Field{V, S}) where {V, S} = local_geometry_type(S)
 
 ClimaComms.context(field::Field) = ClimaComms.context(axes(field))

test/Fields/unit_field.jl

@@ -1118,6 +1118,34 @@ end
     @. f += 1
 end
 
+@testset "Boolean fields" begin
+    FT = Float32
+    space = ExtrudedCubedSphereSpace(;
+        z_elem = 10,
+        z_min = 0,
+        z_max = 1,
+        radius = 10,
+        h_elem = 10,
+        n_quad_points = 4,
+        staggering = Grids.CellCenter(),
+    )
+    bf = Fields.Field(Bool, space)
+    @. bf = true
+    @test all(x -> x == true, Array(parent(bf)))
+    @. bf = 1
+    @test all(x -> x == true, Array(parent(bf)))
+    @. bf = 0
+    @test all(x -> x == false, Array(parent(bf)))
+    @. bf = 0 + bf # test copyto!(bf, ::Braodcasted)
+    @test all(x -> x == false, Array(parent(bf)))
+    if ClimaComms.device() isa ClimaComms.AbstractCPUDevice
+        @test_throws InexactError begin
+            @. bf = 2.0 # no error on gpu
+        end
+    end
+    bf_new = @. bf # test copy()
+end
+
 include("unit_field_multi_broadcast_fusion.jl")
 
 nothing