Fix Integral reductions on immersed grids (#3949)

src/Fields/field.jl

@@ -40,7 +40,7 @@ function validate_field_data(loc, data, grid, indices)
     Fx, Fy, Fz = total_size(grid, loc, indices)
 
     if size(data) != (Fx, Fy, Fz)
-        LX, LY, LZ = loc    
+        LX, LY, LZ = loc
         e = "Cannot construct field at ($LX, $LY, $LZ) with size(data)=$(size(data)). " *
             "`data` must have size ($Fx, $Fy, $Fz)."
         throw(ArgumentError(e))
@@ -179,7 +179,7 @@ function Field(loc::Tuple,
 
     return Field(loc, grid, data, boundary_conditions, indices, operand, status)
 end
-    
+
 Field(z::ZeroField; kw...) = z
 Field(f::Field; indices=f.indices) = view(f, indices...) # hmm...
 
@@ -394,7 +394,7 @@ Return a view of `f` that excludes halo points.
 interior(f::Field) = interior(f.data, location(f), f.grid, f.indices)
 interior(a::OffsetArray, loc, grid, indices) = interior(a, loc, topology(grid), size(grid), halo_size(grid), indices)
 interior(f::Field, I...) = view(interior(f), I...)
-    
+
 # Don't use axes(f) to checkbounds; use axes(f.data)
 Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...)
 
@@ -419,8 +419,8 @@ total_size(f::Field) = total_size(f.grid, location(f), f.indices)
 ##### Move Fields between architectures
 #####
 
-on_architecture(arch, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = 
-    Field{LX, LY, LZ}(on_architecture(arch, field.grid), 
+on_architecture(arch, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} =
+    Field{LX, LY, LZ}(on_architecture(arch, field.grid),
                       on_architecture(arch, field.data),
                       on_architecture(arch, field.boundary_conditions),
                       on_architecture(arch, field.indices),
@@ -513,14 +513,14 @@ const ReducedField = Union{XReducedField,
                            XYReducedField,
                            XYZReducedField}
 
-reduced_dimensions(field::Field)   = ()
-reduced_dimensions(field::XReducedField)   = tuple(1)
-reduced_dimensions(field::YReducedField)   = tuple(2)
-reduced_dimensions(field::ZReducedField)   = tuple(3)
-reduced_dimensions(field::YZReducedField)  = (2, 3)
-reduced_dimensions(field::XZReducedField)  = (1, 3)
-reduced_dimensions(field::XYReducedField)  = (1, 2)
-reduced_dimensions(field::XYZReducedField) = (1, 2, 3)
+reduced_dimensions(::Field)   = ()
+reduced_dimensions(::XReducedField)   = tuple(1)
+reduced_dimensions(::YReducedField)   = tuple(2)
+reduced_dimensions(::ZReducedField)   = tuple(3)
+reduced_dimensions(::YZReducedField)  = (2, 3)
+reduced_dimensions(::XZReducedField)  = (1, 3)
+reduced_dimensions(::XYReducedField)  = (1, 2)
+reduced_dimensions(::XYZReducedField) = (1, 2, 3)
 
 @propagate_inbounds Base.getindex(r::XReducedField, i, j, k) = getindex(r.data, 1, j, k)
 @propagate_inbounds Base.getindex(r::YReducedField, i, j, k) = getindex(r.data, i, 1, k)
@@ -590,6 +590,15 @@ const ReducedAbstractField = Union{XReducedAbstractField,
                                    XYReducedAbstractField,
                                    XYZReducedAbstractField}
 
+reduced_dimensions(::AbstractField)   = ()
+reduced_dimensions(::XReducedAbstractField)   = tuple(1)
+reduced_dimensions(::YReducedAbstractField)   = tuple(2)
+reduced_dimensions(::ZReducedAbstractField)   = tuple(3)
+reduced_dimensions(::YZReducedAbstractField)  = (2, 3)
+reduced_dimensions(::XZReducedAbstractField)  = (1, 3)
+reduced_dimensions(::XYReducedAbstractField)  = (1, 2)
+reduced_dimensions(::XYZReducedAbstractField) = (1, 2, 3)
+
 # TODO: needs test
 Statistics.dot(a::Field, b::Field) = mapreduce((x, y) -> x * y, +, interior(a), interior(b))
 
@@ -605,7 +614,7 @@ const AnyReduction     = typeof(Base.any!)
 initialize_reduced_field!(::SumReduction,     f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, Base.add_sum, true, interior(c))
 initialize_reduced_field!(::ProdReduction,    f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, Base.mul_prod, true, interior(c))
 initialize_reduced_field!(::AllReduction,     f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, &, true, interior(c))
-initialize_reduced_field!(::AnyReduction,     f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, |, true, interior(c))             
+initialize_reduced_field!(::AnyReduction,     f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, |, true, interior(c))
 initialize_reduced_field!(::MaximumReduction, f, r::ReducedAbstractField, c) = Base.mapfirst!(f, interior(r), interior(c))
 initialize_reduced_field!(::MinimumReduction, f, r::ReducedAbstractField, c) = Base.mapfirst!(f, interior(r), interior(c))
 
@@ -658,7 +667,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
     reduction! = Symbol(reduction, '!')
 
     @eval begin
-        
+
         # In-place
         function Base.$(reduction!)(f::Function,
                                     r::ReducedAbstractField,
@@ -710,7 +719,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
     end
 end
 
-function Statistics._mean(f, c::AbstractField, ::Colon; condition = nothing, mask = 0) 
+function Statistics._mean(f, c::AbstractField, ::Colon; condition = nothing, mask = 0)
     operator = condition_operand(f, c, condition, mask)
     return sum(operator) / conditional_length(operator)
 end

src/Operators/Operators.jl

@@ -10,6 +10,10 @@ export Axᶠᶠᶠ, Axᶠᶠᶜ, Axᶠᶜᶠ, Axᶠᶜᶜ, Axᶜᶠᶠ, Axᶜᶠ
 export Ayᶠᶠᶠ, Ayᶠᶠᶜ, Ayᶠᶜᶠ, Ayᶠᶜᶜ, Ayᶜᶠᶠ, Ayᶜᶠᶜ, Ayᶜᶜᶠ, Ayᶜᶜᶜ
 export Azᶠᶠᶠ, Azᶠᶠᶜ, Azᶠᶜᶠ, Azᶠᶜᶜ, Azᶜᶠᶠ, Azᶜᶠᶜ, Azᶜᶜᶠ, Azᶜᶜᶜ
 
+export Axᵃᶜᶜ, Axᵃᶠᶠ, Axᶜᵃᶜ, Axᶠᵃᶠ, Axᶜᶜᵃ, Axᶠᶠᵃ
+export Ayᵃᶜᶜ, Ayᵃᶠᶠ, Ayᶜᵃᶜ, Ayᶠᵃᶠ, Ayᶜᶜᵃ, Ayᶠᶠᵃ
+export Azᵃᶜᶜ, Azᵃᶠᶠ, Azᶜᵃᶜ, Azᶠᵃᶠ, Azᶜᶜᵃ, Azᶠᶠᵃ
+
 # Volumes
 export Vᶠᶠᶠ, Vᶠᶠᶜ, Vᶠᶜᶠ, Vᶠᶜᶜ, Vᶜᶠᶠ, Vᶜᶠᶜ, Vᶜᶜᶠ, Vᶜᶜᶜ
 

test/test_field_scans.jl

@@ -92,7 +92,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k))
                 for T′ in (Tx, Txy)
                     @test T′.operand.operand === T
                 end
-                
+
                 for w′ in (wx, wxy)
                     @test w′.operand.operand === w
                 end
@@ -255,7 +255,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k))
                 push!(results, all(interior(C) .== 1))
             end
 
-            @test mean(results) == 1.0              
+            @test mean(results) == 1.0
         end
 
         @testset "Allocating reductions [$arch_str]" begin
@@ -331,7 +331,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k))
         @testset "Immersed Fields reduction [$(typeof(arch))]" begin
             @info "  Testing reductions of immersed Fields [$(typeof(arch))]"
             underlying_grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(1, 1, 1))
-            
+
             grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> y < 0.5 ? - 0.6 : 0))
             c = Field((Center, Center, Nothing), grid)
 
@@ -363,6 +363,23 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k))
             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
+
+            # See: https://github.com/CliMA/Oceananigans.jl/issues/3948
+            underlying_grid = RectilinearGrid(
+                topology=(Periodic, Periodic, Periodic),
+                size=(3, 3, 3),
+                x=(0, 1), y=(0, 1), z=(0, 1)
+            )
+
+            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> x + y))
+
+            c = CenterField(grid)
+            set!(c, (x, y, z) -> x + y + z)
+
+            max_c² = Field(Reduction(maximum, c^2, dims=3))
+            ∫max_c² = Integral(max_c², dims=(1, 2))
+            compute!(∫max_c²)
+            @test ∫max_c² isa Reduction
         end
     end
 end