Introduces SumOfFields

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.91.4"
+version = "0.91.5"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/ImmersedBoundaries/mask_immersed_field.jl

@@ -1,6 +1,7 @@
 using KernelAbstractions: @kernel, @index
 using Statistics
 using Oceananigans.Fields: location, ZReducedField, Field
+using Oceananigans.Utils: SumOfArrays
 
 instantiate(T::Type) = T()
 instantiate(t) = t
@@ -9,6 +10,13 @@ mask_immersed_field!(field, grid, loc, value) = nothing
 mask_immersed_field!(field::Field, value=zero(eltype(field.grid))) =
     mask_immersed_field!(field, field.grid, location(field), value)
 
+function mask_immersed_field!(sumofarrays::SumOfArrays, value=zero(eltype(sumofarrays.arrays[1])))
+    loc = @inbounds location(sumofarrays.arrays[1])
+    grid = @inbounds sumofarrays.arrays[1].grid
+
+   return mask_immersed_field!(sumofarrays, grid, loc, value)
+end
+
 """
     mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
 
@@ -21,6 +29,17 @@ function mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, val
     return nothing
 end
 
+function mask_immersed_field!(sumofarrays::SumOfArrays, grid::ImmersedBoundaryGrid, loc, value; k, mask)
+    arch = @inbounds architecture(sumofarrays.arrays[1])
+    loc = instantiate.(loc)
+
+    for field in sumofarrays.arrays
+        launch!(arch, grid, :xyz,
+                _mask_immersed_field!, field, loc, grid, value, k, mask)
+    end
+
+    return nothing
+end
 
 @kernel function _mask_immersed_field!(field, loc, grid, value)
     i, j, k = @index(Global, NTuple)
@@ -49,6 +68,18 @@ end
     @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field, mask)
 end
 
+function mask_immersed_field_xy!(sumofarrays::SumOfArrays, grid::ImmersedBoundaryGrid, loc, value; k, mask)
+    arch = @inbounds architecture(sumofarrays.arrays[1])
+    loc = instantiate.(loc)
+
+    for field in sumofarrays.arrays
+        launch!(arch, grid, :xy,
+                _mask_immersed_field_xy!, field, loc, grid, value, k, mask)
+    end
+
+    return nothing
+end
+
 #####
 ##### Masking for GridFittedBoundary
 #####