Updates to masked spaces (#2235)

Author: charleskawczynski

src/DataLayouts/DataLayouts.jl

@@ -2330,8 +2330,17 @@ function ColumnMask(
     @assert Nq isa Integer
     @assert Nh isa Integer
     T = horizontal_layout_type
-    is_active = replace_basetype(T{FT, Nq}(Array{FT}, Nh), Bool)
-    Nijh = Nq * Nq * Nh
+    is_active = replace_basetype(T{FT, Nq}(DA{FT}, Nh), Bool)
+    parent(is_active) .= true
+    return IJHMask(is_active)
+end
+
+union_all_type(::Type{T}) where {T} = T.name.wrapper
+
+function IJHMask(is_active::Union{IJFH, IJHF})
+    DA = union_all_type(typeof(parent(is_active)))
+    (Ni, Nj, _, _, Nh) = size(is_active)
+    Nijh = Ni * Nj * Nh
     i_map = zeros(Int, Nijh)
     j_map = zeros(Int, Nijh)
     h_map = zeros(Int, Nijh)

src/Fields/broadcast.jl

@@ -92,7 +92,7 @@ function Base.copy(
         throw(BroadcastInferenceError(bc))
     end
     # We can trust it and defer to the simpler `copyto!`
-    return copyto!(similar(bc, ElType), bc, DataLayouts.NoMask())
+    return copyto!(similar(bc, ElType), bc, Spaces.get_mask(axes(bc)))
 end
 
 Base.@propagate_inbounds function slab(

src/Spaces/spectralelement.jl

@@ -20,6 +20,11 @@ function Base.show(io::IO, space::AbstractSpectralElementSpace)
     indent = get(io, :indent, 0)
     iio = IOContext(io, :indent => indent + 2)
     println(io, nameof(typeof(space)), ":")
+    if get_mask(space) isa DataLayouts.NoMask
+        println(iio, " "^(indent + 2), "mask_enabled: false")
+    else
+        println(iio, " "^(indent + 2), "mask_enabled: true")
+    end
     if hasfield(typeof(grid(space)), :topology)
         # some reduced spaces (like slab space) do not have topology
         print(iio, " "^(indent + 2), "context: ")

test/Spaces/unit_spaces.jl

@@ -23,11 +23,15 @@ import ClimaCore:
     DeviceSideDevice
 
 using ClimaCore.CommonSpaces
+using ClimaCore.Utilities.Cache
 import ClimaCore.DataLayouts: IJFH, VF, slab_index
 
 on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice
 
 @testset "2D spaces with mask" begin
+    # We need to test a fresh instance of the spaces, since
+    # masked spaces include data set by users.
+    Cache.clean_cache!()
     FT = Float64
     context = ClimaComms.context()
     x_max = FT(1)
@@ -57,6 +61,7 @@ on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice
     hspace = Spaces.SpectralElementSpace2D(htopology, quad; enable_mask = true)
     mask = Spaces.get_mask(hspace)
     @test mask isa DataLayouts.IJHMask
+    @test all(x -> x == true, parent(mask.is_active)) # test that default is true
     Spaces.set_mask!(hspace) do coords
         coords.x > 0.5
     end
@@ -71,6 +76,9 @@ on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice
     @. f = 1 + ᶜx * 0 # tests copyto!
     @test count(iszero, parent(f)) == 2
 
+    fbc = @. 1 + ᶜx * 0 # tests copy
+    @test Spaces.get_mask(axes(fbc)) isa DataLayouts.IJHMask
+
     FT = Float64
     ᶜspace = ExtrudedCubedSphereSpace(
         FT;
@@ -92,8 +100,7 @@ on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice
     # Test that mask-field assignment works:
     # TODO: we should make this easier
     is_active = similar(mask.is_active)
-    _is_active = Fields.field_values(float.(hᶠcoords.lat .> 0.5))
-    is_active .= DataLayouts.replace_basetype(_is_active, Bool)
+    parent(is_active) .= parent(hᶠcoords.lat) .> 0.5
     Spaces.set_mask!(ᶜspace, is_active)
 
     Spaces.set_mask!(ᶜspace) do coords
@@ -131,6 +138,8 @@ on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice
     ᶠspace_no_mask = Spaces.face_space(ᶜspace_no_mask)
     ᶠcoords_no_mask = Fields.coordinate_field(ᶠspace_no_mask)
     c_no_mask = Fields.Field(FT, ᶜspace_no_mask)
+    @test_throws ErrorException("Broacasted spaces are not the same.") @. c_no_mask +
+                                                                          ᶜf
     ᶠf_no_mask = Fields.Field(FT, ᶠspace_no_mask)
     if ClimaComms.device(ᶜspace_no_mask) isa ClimaComms.CUDADevice
         @. c_no_mask = div(Geometry.WVector(foo(ᶠf_no_mask, ᶠcoords_no_mask)))
@@ -162,6 +171,7 @@ end
 
     expected_repr = """
     SpectralElementSpace1D:
+      mask_enabled: false
       context: SingletonCommsContext using $(nameof(typeof(device)))
       mesh: 1-element IntervalMesh of IntervalDomain: x ∈ [-3.0,5.0] (periodic)
       quadrature: 4-point Gauss-Legendre-Lobatto quadrature"""
@@ -348,6 +358,7 @@ end
     space = Spaces.SpectralElementSpace2D(grid_topology, quad)
     @test repr(space) == """
     SpectralElementSpace2D:
+      mask_enabled: false
       context: SingletonCommsContext using CPUSingleThreaded
       mesh: 1×1-element RectilinearMesh of RectangleDomain: x ∈ [-3.0,5.0] (periodic) × y ∈ [-2.0,8.0] (:south, :north)
       quadrature: 4-point Gauss-Legendre-Lobatto quadrature"""