extending the split explicit free surface to behave correctly

src/split_explicit_free_surface.jl

@@ -0,0 +1,49 @@
+using Oceananigans.Grids: halo_size, with_halo
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitState, SplitExplicitFreeSurface, calculate_column_height! 
+
+import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface, SplitExplicitAuxiliaryFields
+import Oceananigans.Models.HydrostaticFreeSurfaceModels: augmented_kernel_offsets, augmented_kernel_size
+
+function SplitExplicitAuxiliaryFields(grid::TRG)
+
+    Gᵁ = Field((Face,   Center, Nothing), grid)
+    Gⱽ = Field((Center, Face,   Nothing), grid)
+
+    Hᶠᶜ = Field((Face,   Center, Nothing), grid)
+    Hᶜᶠ = Field((Center, Face,   Nothing), grid)
+
+    calculate_column_height!(Hᶠᶜ, (Face, Center, Center))
+    calculate_column_height!(Hᶜᶠ, (Center, Face, Center))
+
+    fill_halo_regions!((Hᶠᶜ, Hᶜᶠ))
+
+    Nx, Ny, _ = size(grid)
+    Hx, Hy, _ = halo_size(grid)
+
+    kernel_parameters = (Nx, Ny + Hy - 1)
+
+    return SplitExplicitAuxiliaryFields(Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, kernel_parameters)
+end
+
+# Internal function for HydrostaticFreeSurfaceModel
+function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::TRG)
+
+        settings  = free_surface.settings 
+
+        old_halos  = halo_size(grid)
+        Nsubsteps  = length(settings.substepping.averaging_weights)
+
+        extended_halos = tripolar_split_explicit_halos(old_halos, Nsubsteps+1, grid)         
+        extended_grid  = with_halo(extended_halos, grid)
+
+        Nze = size(extended_grid, 3)
+        η = ZFaceField(extended_grid, indices = (:, :, Nze+1))
+
+        return SplitExplicitFreeSurface(η,
+                                        SplitExplicitState(extended_grid, settings.timestepper),
+                                        SplitExplicitAuxiliaryFields(extended_grid),
+                                        free_surface.gravitational_acceleration,
+                                        free_surface.settings)
+end
+
+@inline tripolar_split_explicit_halos(old_halos, step_halo, grid) = old_halos[1], max(step_halo, old_halos[2])

src/with_halo.jl

@@ -0,0 +1,17 @@
+using Oceananigans.Grids: architecture, cpu_face_constructor_z
+
+import Oceananigans.Grids: with_halo
+
+function with_halo(new_halo, old_grid::TripolarGrid)
+
+    size = (old_grid.Nx, old_grid.Ny, old_grid.Nz)
+
+    z = cpu_face_constructor_z(old_grid)
+
+    new_grid = TripolarGrid(architecture(old_grid), eltype(old_grid);
+                                            size, z, 
+                                            radius = old_grid.radius,
+                                            halo = new_halo)
+
+    return new_grid
+end