Yet another ZStar implementation

[simone-silvestri]

Refactoring

Following #3411 it is clear that the grids require some refactor to allow ZStar.
This PR implements another proposal for ZStar which hinges on changing the grids to have just one z field.
The case of a static z (the only allowed in Oceananigans at the moment) is covered with a StaticVerticalCoordinate type that contains all the specifications that were in the grid before:

Oceananigans.jl/src/Grids/vertical_coordinate.jl

Lines 14 to 19 in 9e05e21

	struct StaticVerticalCoordinate{C, D} <: AbstractVerticalCoordinate
	cᶠ :: C
	cᶜ :: C
	Δᶜ :: D
	Δᶠ :: D
	end

This refactor should change nothing in the user interface nor the internals of the code, since the whole API should depend on znode and Δz, but just reorganizes the vertical coordinate in a compact type that can be modified.

ZStar Implementation

With this refactor, a ZStar coordinate is implemented through a ZStarVerticalCoordinate type

Oceananigans.jl/src/Grids/vertical_coordinate.jl

Lines 36 to 48 in 9e05e21

	struct ZStarVerticalCoordinate{C, D, E, CC, FC, CF, FF} <: AbstractVerticalCoordinate
	cᶠ :: C
	cᶜ :: C
	Δᶜ :: D
	Δᶠ :: D
	ηⁿ :: E
	e₃ᶜᶜⁿ :: CC
	e₃ᶠᶜⁿ :: FC
	e₃ᶜᶠⁿ :: CF
	e₃ᶠᶠⁿ :: FF
	e₃ᶜᶜ⁻ :: CC
	∂t_e₃ :: CC
	end

e₃ is the vertical scaling of the grid spacing defined as

$$\frac{\partial z}{\partial r} = \frac{H + \eta}{H}$$

with $z$ the actual spatial coordinate (moving with the free surface) and $r$ is the reference vertical coordinate also called $z^\star$ (equivalent to a static coordinate, or $z$ when $\eta = 0$).
ηⁿ is the value of the free surface at the current time step, needed to calculate znodes

$$z = r \cdot e_3 + \eta \ .$$

Finally $H$ is the output of static_column_depth.
As a consequence of the definitions, all $z$ properties (znodes, zspacings, Δz, etc...) are replaced with $r$ (rnodes, rspacings,Δr, etc...) and the $z$ properties for a zstar vertical coordinate are calculated appropriately in the Operators module (https://github.com/CliMA/Oceananigans.jl/blob/ss/new-zstar/src/Operators/variable_grid_operators.jl)

An example of the proposed user interface is shown in the validation/z_star_coordinate examples

Changes in the model timestepping

Another requirement is changing the internals of the HydrostaticFreeSurfaceModel to solve the correct equations, in particular, the notable changes are:

(1) Updating the grid at each timestep in the update_state! step

(2) including the gradient of the grid in the momentum equations, calculated as:

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl

Lines 85 to 100 in 9e05e21

	#####
	##### ZStar-specific implementation of the additional terms to be included in the momentum equations
	#####
	@inline ∂x_z(i, j, k, grid) = @inbounds ∂xᶠᶜᶜ(i, j, k, grid, znode, Center(), Center(), Center())
	@inline ∂y_z(i, j, k, grid) = @inbounds ∂yᶜᶠᶜ(i, j, k, grid, znode, Center(), Center(), Center())
	@inline grid_slope_contribution_x(i, j, k, grid::ZStarGridOfSomeKind, ::Nothing, model_fields) = zero(grid)
	@inline grid_slope_contribution_y(i, j, k, grid::ZStarGridOfSomeKind, ::Nothing, model_fields) = zero(grid)
	@inline grid_slope_contribution_x(i, j, k, grid::ZStarGridOfSomeKind, buoyancy, model_fields) =
	ℑxᶠᵃᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.model, model_fields) * ∂x_z(i, j, k, grid)
	@inline grid_slope_contribution_y(i, j, k, grid::ZStarGridOfSomeKind, buoyancy, model_fields) =
	ℑyᵃᶠᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.model, model_fields) * ∂y_z(i, j, k, grid)

(3) Changing the computation of the vertical velocity to include the movement of the grid

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl

Lines 22 to 35 in 9e05e21

	@kernel function _compute_w_from_continuity!(U, grid)
	i, j = @index(Global, NTuple)
	@inbounds U.w[i, j, 1] = 0
	for k in 2:grid.Nz+1
	δh_u = flux_div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) / Azᶜᶜᶜ(i, j, k-1, grid)
	∂t_s = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_e₃(i, j, k-1, grid)
	immersed = immersed_cell(i, j, k-1, grid)
	Δw = δh_u + ifelse(immersed, 0, ∂t_s) # We do not account for grid changes in immersed cells
	@inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δw
	end
	end

(4) Advancing $e_3\theta$ instead of $\theta$ in the tracer equations and subsequently unscale it back after the grid scaling factor at the new time step $e_3^{n+1}$ is known

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

Lines 111 to 128 in 9e05e21

	@kernel function _ab2_step_tracer_field!(θ, grid, Δt, χ, Gⁿ, G⁻)
	i, j, k = @index(Global, NTuple)
	FT = eltype(χ)
	C₁ = convert(FT, 1.5) + χ
	C₂ = convert(FT, 0.5) + χ
	eⁿ = e₃ⁿ(i, j, k, grid, Center(), Center(), Center())
	e⁻ = e₃⁻(i, j, k, grid, Center(), Center(), Center())
	@inbounds begin
	∂t_sθ = C₁ * eⁿ * Gⁿ[i, j, k] - C₂ * e⁻ * G⁻[i, j, k]
	# We store temporarily sθ in θ. the unscaled θ will be retrived later on with `unscale_tracers!`
	θ[i, j, k] = eⁿ * θ[i, j, k] + convert(FT, Δt) * ∂t_sθ
	end
	end

(5) Adding a non-linear free surface by changing the static_column_depth in the split explicit free surface implementation to a dynamic_column_depth defined as

Oceananigans.jl/src/ImmersedBoundaries/zstar_immersed_grid.jl

Lines 14 to 17 in 9e05e21

	@inline dynamic_column_depthᶜᶜᵃ(i, j, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶜᶜᵃ(i, j, grid) + η[i, j, grid.Nz+1]
	@inline dynamic_column_depthᶜᶠᵃ(i, j, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶜᶠᵃ(i, j, grid) + ℑxᶠᵃᵃ(i, j, grid.Nz+1, grid, η)
	@inline dynamic_column_depthᶠᶜᵃ(i, j, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶠᶜᵃ(i, j, grid) + ℑyᵃᶠᵃ(i, j, grid.Nz+1, grid, η)
	@inline dynamic_column_depthᶠᶠᵃ(i, j, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶠᶠᵃ(i, j, grid) + ℑxyᶠᶠᵃ(i, j, grid.Nz+1, grid, η)

Note: these changes are valid only for the QuasiAdamsBashort2 timestepper. Support for the SplitRungeKutta3 timestepper requires a bit more work and can be done when all the infrastructure is in place

OutputWriters?

The last piece of the puzzle (not implemented in this PR) would be to change the OutputWriters to include by default the variable grid properties in the timeseries field of the jld2 / netcdf writer.

Possibility for improvements

I got a bit carried away and completed a working implementation to make sure it could be done this way, but I am happy to change just about everything in here. I would like to know what people think about this implementation and what people suggest especially in terms of
(1) Implementation
(2) Variable naming

I can also split this PR in a couple of ones, probably one that includes the refactor to the grids (implementation of a StaticVerticalCoordinate) and one that implements a working version of a ZStarVerticalCoordinate.