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.

[ali-ramadhan]

Super exciting PR! Would love to try z* in some steep topography setups.

I'm not super familiar with z* coordinates but some thoughts:

1. I do like the modularity you've introduced with StaticVerticalCoordinate and ZStarVerticalCoordinate! Looks like it'll be easy to switch between the two.
2. Why use $r$ (rnodes, rspacings, etc.)? I ask because $r$ suggests radial coordinates on a sphere where the horizontal spacing of each grid cell increases as $r$ increases. But as far as I know from looking at Figure 1 of Adcroft & Campin (2004) [see below] the vertical grid lines are still straight in z* coordinates.
3. More a question of notation: I know you've defined e₃ is the vertical scaling of the grid spacing $\partial z / \partial r$. Is this the unit vector in the third dimension on the new $z^\star$ grid? I guess $e_n$ usually denotes the unit vector in the nth dimension (or maybe $\hat{e}_n$).

Would it be more accurate to use $z^\star$ instead of $r$ then $e_3 = \partial z / \partial z^\star$? I could be misunderstanding the formulation!

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

Does this mean Oceananigans now supports a non-linear free surface?

[simone-silvestri]

Thanks for the feedback. I was indeed not super sure about the naming so we can decide how to make it more understandable.

2. Why use $r$ (rnodes, rspacings, etc.)? I ask because $r$ suggests radial coordinates on a sphere where the horizontal spacing of each grid cell increases as $r$ increases. But as far as I know from looking at Figure 1 of Adcroft & Campin (2004) [see below] the vertical grid lines are still straight in z* coordinates.

I have used r for reference. This coordinate stands for a generalized vertical coordinate equivalent to the z in a static geopotential grid (r in a zstar grid is equivalent to the z of a grid with a free surface equal to zero). Adcroft and Campin use r for a generalized coordinate in the appendix (zstar is a particular type of r in this case). But I am down to finding another name to avoid confusion with radius. In NEMO they call it s for example, which might be compact and unique enough.

3. More a question of notation: I know you've defined e₃ is the vertical scaling of the grid spacing $\partial z / \partial r$. Is this the unit vector in the third dimension on the new $z^\star$ grid? I guess $e_n$ usually denotes the unit vector in the nth dimension (or maybe $\hat{e}_n$).

I have chosen e₃ because it is used also in NEMO and it is compact enough. My previous version had vertical_scaling which I found out to be a little too verbose for an operator. We can decide on another symbol we want to use for this quantity. Maybe ∂z? or it might be confused with the derivative in z? If we stick with the r notation ∂r_z would be probably the best choice (or ∂s_z if we go with s)

Does this mean Oceananigans now supports a non-linear free surface?

Yep, with zstar and a splitexplicit free surface, we have a non-linear free surface implementation automatically.

[simone-silvestri]

actually, the more I think about it, probably s fits the characteristics of a new coordinate (unambiguous, compact and general) and e₃ can become ∂s_z. The problem with using z* is that it does not represent a general coordinate.