Fixes error when reducing Fields with a condition on `ImmersedBou…

…ndaryGrids` (#3440) * add missing method * add tests * Update src/ImmersedBoundaries/immersed_reductions.jl Co-authored-by: Gregory L. Wagner <[email protected]> * clearer condition_operand * more spaced out syntax in test * fix typo * remove allowscalar * Update test/test_field_reductions.jl Co-authored-by: Gregory L. Wagner <[email protected]> * Update test/test_field_reductions.jl Co-authored-by: Gregory L. Wagner <[email protected]> * fix naming in test * add method for `::BitArray` --------- Co-authored-by: Gregory L. Wagner <[email protected]> Co-authored-by: Simone Silvestri <[email protected]>
CliMA · Feb 1, 2024 · c3151d6 · c3151d6
1 parent dc34b80
commit c3151d6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/src/Architectures.jl b/src/Architectures.jl
@@ -61,6 +61,9 @@ arch_array(::CPU, a::CuArray) = Array(a)
 arch_array(::GPU, a::Array)   = CuArray(a)
 arch_array(::GPU, a::CuArray) = a
 
+arch_array(::CPU, a::BitArray) = a
+arch_array(::GPU, a::BitArray) = CuArray(a)
+
 arch_array(::GPU, a::SubArray{<:Any, <:Any, <:CuArray}) = a
 arch_array(::CPU, a::SubArray{<:Any, <:Any, <:CuArray}) = Array(a)
 

diff --git a/src/ImmersedBoundaries/immersed_reductions.jl b/src/ImmersedBoundaries/immersed_reductions.jl
@@ -21,6 +21,13 @@ const IF = AbstractField{<:Any, <:Any, <:Any, <:ImmersedBoundaryGrid}
 @inline condition_operand(func::Function,         op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask)
 @inline condition_operand(func::typeof(identity), op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask)
 
+@inline function condition_operand(func::Function, op::IF, cond::AbstractArray, mask)
+    arch = architecture(op.grid)
+    arch_condition = arch_array(arch, cond)
+    ni_condition = NotImmersed(arch_condition)
+    return ConditionalOperation(op; func, condition=ni_condition, mask)
+end
+
 @inline conditional_length(c::IF)       = conditional_length(condition_operand(identity, c, nothing, 0))
 @inline conditional_length(c::IF, dims) = conditional_length(condition_operand(identity, c, nothing, 0), dims)
 

diff --git a/test/test_field_reductions.jl b/test/test_field_reductions.jl
@@ -228,13 +228,33 @@ trilinear(x, y, z) = x + y + z
             c = Field((Center, Center, Nothing), grid)
 
             set!(c, (x, y) -> y)
-            @test maximum(c) == CUDA.@allowscalar grid.yᵃᶜᵃ[1]
+            @test maximum(c) == grid.yᵃᶜᵃ[1]
 
             grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> y < 0.5 ? - 0.6 : -0.4))
             c = Field((Center, Center, Nothing), grid)
 
             set!(c, (x, y) -> y)
-            @test maximum(c) == CUDA.@allowscalar grid.yᵃᶜᵃ[3]
+            @test maximum(c) == grid.yᵃᶜᵃ[3]
+
+            underlying_grid = RectilinearGrid(arch, size = (1, 1, 8), extent=(1, 1, 1))
+
+            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> -3/4))
+            c = Field((Center, Center, Center), grid)
+
+            set!(c, (x, y, z) -> -z)
+            @test maximum(c) == Array(interior(c))[1, 1, 3]
+            c_condition = interior(c) .< 0.5
+            avg_c_smaller_than_½ = Array(interior(compute!(Field(Average(c, condition=c_condition)))))
+            @test avg_c_smaller_than_½[1, 1, 1] == 0.25
+
+            zᶜᶜᶜ = KernelFunctionOperation{Center, Center, Center}(znode, grid, Center(), Center(), Center())
+            ci = Array(interior(c)) # transfer to CPU
+            bottom_half_average_manual = (ci[1, 1, 3] + ci[1, 1, 4]) / 2
+            bottom_half_average = Average(c; condition=(zᶜᶜᶜ .< -1/2))
+            bottom_half_average_field = Field(bottom_half_average)
+            compute!(bottom_half_average_field)
+            bottom_half_average_array = Array(interior(bottom_half_average_field))
+            @test bottom_half_average_array[1, 1, 1] == bottom_half_average_manual
         end
     end
 end