Add support for ZStar in RK3 (#4156)

* add stuff

* changes

* test it out

* let's go!

* update

* adapt tracer

* try it out

* fix this

* let's go

* add an RK3 test with ZStar

* lock release

Author: simone-silvestri

src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl

@@ -81,7 +81,7 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) =
 
 @kernel function _explicit_rk3_step_free_surface!(η, Δt, γⁿ, ζⁿ, Gⁿ, η⁻, Nz)
     i, j = @index(Global, NTuple)
-    @inbounds η[i, j, Nz+1] += ζⁿ * η⁻[i, j, k] + γⁿ * (η[i, j, k] + convert(FT, Δt) * Gⁿ[i, j, k])
+    @inbounds η[i, j, Nz+1] += ζⁿ * η⁻[i, j, Nz+1] + γⁿ * (η[i, j, Nz+1] + Δt * Gⁿ[i, j, Nz+1])
 end
 
 @kernel function _explicit_ab2_step_free_surface!(η, Δt, χ, Gηⁿ, Gη⁻, Nz)

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

@@ -89,8 +89,10 @@ function ab2_step_tracers!(tracers, model, Δt, χ)
             G⁻ = model.timestepper.G⁻[tracer_name]
             tracer_field = tracers[tracer_name]
             closure = model.closure
+            grid = model.grid
 
-            ab2_step_tracer_field!(tracer_field, model.grid, Δt, χ, Gⁿ, G⁻)
+            FT = eltype(grid)
+            launch!(architecture(grid), grid, :xyz, _ab2_step_tracer_field!, tracer_field, grid, convert(FT, Δt), χ, Gⁿ, G⁻)
 
             implicit_step!(tracer_field,
                            model.timestepper.implicit_solver,
@@ -105,9 +107,6 @@ function ab2_step_tracers!(tracers, model, Δt, χ)
     return nothing
 end
 
-ab2_step_tracer_field!(tracer_field, grid, Δt, χ, Gⁿ, G⁻) =
-    launch!(architecture(grid), grid, :xyz, _ab2_step_tracer_field!, tracer_field, grid, Δt, χ, Gⁿ, G⁻)
-
 #####
 ##### Tracer update in mutable vertical coordinates 
 #####
@@ -129,7 +128,7 @@ ab2_step_tracer_field!(tracer_field, grid, Δt, χ, Gⁿ, G⁻) =
 
         # We store temporarily σθ in θ. 
         # The unscaled θ will be retrieved with `unscale_tracers!`
-        θ[i, j, k] = σᶜᶜⁿ * θ[i, j, k] + convert(FT, Δt) * ∂t_σθ
+        θ[i, j, k] = σᶜᶜⁿ * θ[i, j, k] + Δt * ∂t_σθ
     end
 end
 

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl

@@ -39,5 +39,6 @@ previous_hydrostatic_tendency_fields(::Val{:SplitRungeKutta3}, args...) = nothin
 function previous_hydrostatic_tendency_fields(::Val{:SplitRungeKutta3}, velocities, free_surface::SplitExplicitFreeSurface, args...)
     U = similar(free_surface.barotropic_velocities.U)
     V = similar(free_surface.barotropic_velocities.V)
-    return (; U=U, V=V)
+    η = similar(free_surface.η)
+    return (; U=U, V=V, η=η)
 end

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl

@@ -1,15 +1,15 @@
-using Oceananigans.Fields: location
+using Oceananigans.Fields: location, instantiated_location
 using Oceananigans.TurbulenceClosures: implicit_step!
 using Oceananigans.ImmersedBoundaries: get_active_cells_map, get_active_column_map
 
-import Oceananigans.TimeSteppers: split_rk3_substep!, _split_rk3_substep_field!
+import Oceananigans.TimeSteppers: split_rk3_substep!, _split_rk3_substep_field!, store_fields!
 
 function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζⁿ)
 
     grid         = model.grid
     timestepper  = model.timestepper
     free_surface = model.free_surface
-
+    
     compute_free_surface_tendency!(grid, model, free_surface)
 
     rk3_substep_velocities!(model.velocities, model, Δt, γⁿ, ζⁿ)
@@ -18,8 +18,7 @@ function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζ
     # Full step for Implicit and Split-Explicit, substep for Explicit
     step_free_surface!(free_surface, model, timestepper, Δt)
 
-    # Average free surface variables 
-    # in the second stage
+    # Average free surface variables in the second stage
     if model.clock.stage == 2
         rk3_average_free_surface!(free_surface, grid, timestepper, γⁿ, ζⁿ)
     end
@@ -46,11 +45,17 @@ function rk3_average_free_surface!(free_surface::SplitExplicitFreeSurface, grid,
 
     Uⁿ⁻¹ = timestepper.Ψ⁻.U
     Vⁿ⁻¹ = timestepper.Ψ⁻.V
+    ηⁿ⁻¹ = timestepper.Ψ⁻.η    
     Uⁿ   = free_surface.barotropic_velocities.U
     Vⁿ   = free_surface.barotropic_velocities.V
+    ηⁿ   = free_surface.η 
 
     launch!(arch, grid, :xy, _rk3_average_free_surface!, Uⁿ, grid, Uⁿ⁻¹, γⁿ, ζⁿ)
     launch!(arch, grid, :xy, _rk3_average_free_surface!, Vⁿ, grid, Vⁿ⁻¹, γⁿ, ζⁿ)
+    
+    # Averaging the free surface is only required for a grid with Mutable vertical coordinates, 
+    # which needs to update the grid based on the value of the free surface    
+    launch!(arch, grid, :xy, _rk3_average_free_surface!, ηⁿ, grid, ηⁿ⁻¹, γⁿ, ζⁿ)
 
     return nothing
 end
@@ -97,19 +102,20 @@ function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ)
 
     closure = model.closure
     grid = model.grid
+    FT = eltype(grid)
 
     # Tracer update kernels
     for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
 
         Gⁿ = model.timestepper.Gⁿ[tracer_name]
         Ψ⁻ = model.timestepper.Ψ⁻[tracer_name]
-        tracer_field = tracers[tracer_name]
+        θ  = tracers[tracer_name]
         closure = model.closure
 
         launch!(architecture(grid), grid, :xyz,
-                _split_rk3_substep_field!, tracer_field, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻)
+                _split_rk3_substep_tracer_field!, θ, grid, convert(FT, Δt), γⁿ, ζⁿ, Gⁿ, Ψ⁻)
 
-        implicit_step!(tracer_field,
+        implicit_step!(θ,
                        model.timestepper.implicit_solver,
                        closure,
                        model.diffusivity_fields,
@@ -119,4 +125,56 @@ function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ)
     end
 
     return nothing
-end
+end
+
+#####
+##### Tracer update in mutable vertical coordinates 
+#####
+
+# σθ is the evolved quantity. 
+# We store temporarily σθ in θ. Once σⁿ⁺¹ is known we can retrieve θⁿ⁺¹ 
+# with the `unscale_tracers!` function. Ψ⁻ is the previous tracer already scaled
+# by the vertical coordinate scaling factor: ψ⁻ = σ * θ
+@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) 
+    i, j, k = @index(Global, NTuple)
+
+    σᶜᶜⁿ = σⁿ(i, j, k, grid, Center(), Center(), Center())
+    @inbounds θ[i, j, k] = ζⁿ * Ψ⁻[i, j, k] + γⁿ * σᶜᶜⁿ * (θ[i, j, k] + Δt * Gⁿ[i, j, k])
+end
+
+# We store temporarily σθ in θ. 
+# The unscaled θ will be retrieved with `unscale_tracers!`
+@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, ::Nothing, ::Nothing, Gⁿ, Ψ⁻) 
+    i, j, k = @index(Global, NTuple)
+    @inbounds θ[i, j, k] = Ψ⁻[i, j, k] + Δt * Gⁿ[i, j, k] * σⁿ(i, j, k, grid, Center(), Center(), Center())
+end
+
+##### 
+##### Storing previous fields for the RK3 update
+##### 
+
+# Tracers are multiplied by the vertical coordinate scaling factor
+@kernel function _store_tracer_fields!(Ψ⁻, grid, Ψⁿ) 
+    i, j, k = @index(Global, NTuple)
+    @inbounds Ψ⁻[i, j, k] = Ψⁿ[i, j, k] * σⁿ(i, j, k, grid, Center(), Center(), Center())
+end
+
+function store_fields!(model::HydrostaticFreeSurfaceModel)
+    
+    previous_fields = model.timestepper.Ψ⁻
+    model_fields = prognostic_fields(model)
+    grid = model.grid
+    arch = architecture(grid)
+
+    for name in keys(model_fields)
+        Ψ⁻ = previous_fields[name]
+        Ψⁿ = model_fields[name]
+        if name ∈ keys(model.tracers) # Tracers are stored with the grid scaling
+            launch!(arch, grid, :xyz, _store_tracer_fields!, Ψ⁻, grid, Ψⁿ)
+        else # Velocities and free surface are stored without the grid scaling
+            parent(Ψ⁻) .= parent(Ψⁿ) 
+        end
+    end
+
+    return nothing
+end

test/test_zstar_coordinate.jl

@@ -197,5 +197,31 @@ end
                 end
             end
         end
+
+        @info "  Testing a ZStar and Runge Kutta 3rd order time stepping"
+
+        topology = topologies[2]
+        rtg  = RectilinearGrid(arch; size = (10, 10, 20), x = (0, 100kilometers), y = (-10kilometers, 10kilometers), topology, z = z_uniform)
+        llg  = LatitudeLongitudeGrid(arch; size = (10, 10, 20), latitude = (0, 1), longitude = (0, 1), topology, z = z_uniform)
+        irtg = ImmersedBoundaryGrid(rtg,   GridFittedBottom((x, y) -> rand() - 10))
+        illg = ImmersedBoundaryGrid(llg,   GridFittedBottom((x, y) -> rand() - 10))
+
+        for grid in [rtg, llg, irtg, illg]
+
+            split_free_surface = SplitExplicitFreeSurface(grid; cfl = 0.75)
+            model = HydrostaticFreeSurfaceModel(; grid, 
+                                                free_surface = split_free_surface, 
+                                                tracers = (:b, :c), 
+                                                timestepper = :SplitRungeKutta3,
+                                                buoyancy = BuoyancyTracer(),
+                                                vertical_coordinate = ZStar())
+
+            bᵢ(x, y, z) = x < grid.Lx / 2 ? 0.06 : 0.01 
+
+            set!(model, c = (x, y, z) -> rand(), b = bᵢ)
+
+            Δt = 2minutes
+            test_zstar_coordinate(model, 100, Δt)
+        end
     end
 end

validation/distributed_simulations/distributed_geostrophic_adjustment.jl

@@ -21,7 +21,7 @@ topo = (Bounded, Periodic, Bounded)
 
 partition = Partition(x = Sizes(10, 13, 18, 39))
 
-arch = Distributed(CPU(); partition)
+arch = Distributed(CPU())
 
 # Distribute problem irregularly
 Nx = 80
@@ -46,6 +46,7 @@ coriolis = FPlane(f=1e-4)
 
 model = HydrostaticFreeSurfaceModel(; grid,
                                       coriolis,
+                                      timestepper = :SplitRungeKutta3,
                                       free_surface = SplitExplicitFreeSurface(grid; substeps=10))
 
 gaussian(x, L) = exp(-x^2 / 2L^2)