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.

[NoraLoose]

@simone-silvestri @glwagner thanks for clarifying the use of the wrapper functions to set up the z-star coordinate.

To finish the discussion point on the drift in tracer conservation, I ran the experiment not just for 1 year, but for 10 years. The conclusions are:

• Blue line: We will keep losing the tracer (in a sign-definite fashion), but not at the same rate. While the relative error increases from 1e-16 to 1e-13 over the course of the first year (this is the result we discussed above), it only further increases to 1e-12 over the course of the next 9 years.

Maybe it plays a role that I'm releasing the tracer into a surface cell, and over time it sinks, so it will hit more cells in the vertical. In other words, the effective Nz grows over time.

• Orange line: To test the hypothesis that the magnitude of the error depends on the effective number of grid cells that the tracer takes up, I released a tracer in all surface cells (orange line), not just one surface cell as shown by the blue line. The error is larger initially (which I guess confirms the hypothesis), but after ~1.5 years becomes smaller than the error implied by the single-cell release (which is not really in line with the hypothesis). Not sure what conclusions to draw from that...

TLDR: I think we can feel good about the tracer conservation of the z-star coordinate.

[glwagner]

@NoraLoose another test could release the tracer at depth, then we can see if there is a specific issue with the surface or not.

I sort of agree that this is completely negligible compared to a real signal so for all intents and purposes tracers are conserved. But we should still try to wrap our heads around this eventually I think.

[glwagner]

what happens if the grid is stationary

[glwagner]

confused about this. It makes it look like we always have a moving grid

[glwagner]

what metric?

[glwagner]

I don't think this function does what the comment says. A moving coordinate is presumably implemented by a model. The comment should be updated to state exactly what this function does. PResumably, it makes it possible to have a moving coordinate, somehow.

[glwagner]

Why does this return two things

[glwagner]

Isn't this more general than z*? It's just a vertical coordinate that can evolve in time

[glwagner]

I don't think any code in Grids has to do with a free surface. It must be the case that this code is more general than simply something that follows the free surface. I'm guessing that this is simply an abstraction for a vertical coordinate that can change in time, and it can change according to whatever rule a particular model wants to implement.

[glwagner]

I don't think this has to follow a free surface. That is just a particular choice that is made in another module (HydrostaticFreeSurfaceModels)

[glwagner]

same comment as above

[glwagner]

Suggested change

	@inline dynamic_column_depthᶜᶜᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = static_column_depthᶜᶜᵃ(i, j, grid) + @inbounds η[i, j, k]
	@inline column_depthᶜᶜᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = static_column_depthᶜᶜᵃ(i, j, grid) + @inbounds η[i, j, k]

This is just the "column depth" isn't it? The concept of dynamic vs static doesn't make sense for a fixed grid. Better to refer to this as something that generalizes to all grids. Then we have "static column depth" for the specific case that it's needed.

[glwagner]

why is this 1 for the case that hits the conditional?

[simone-silvestri]

I was thinking that a zero dz might create some NaNs where we don't want them. This assumes that if the height is zero (which basically means that we do not have active cells in the column) we revert to a z-coordinate representation

[glwagner]

I think it's good to define that behavior. Maybe we will learn more about the implications when we have tests. Should we make a comment about it?

[glwagner]

This comment says something about a moving grid, but I think the code here is general and also works for static grids. This comment is very confusing, it will make people think that we always have a moving grid

[glwagner]

This code only works for a HydrostaticFreeSurfaceModel?

[glwagner]

should this be time_variable_grid_operators.jl?

[glwagner]

just "variable" grid is confusing because one might interpret this as referring to a stretched grid.

[glwagner]

what are the results of this test?

[simone-silvestri]

It calculates the drift in a lock-exchange simulation with an sloping immersed boundary.
The drift stays within 3e-16

julia> @show drift
drift = Any[0.0, -1.734723475976807e-16, -2.1510571102112408e-16, -3.469446951953614e-17, -8.326672684688674e-17, -1.0408340855860843e-16, -6.245004513516506e-17, -1.942890293094024e-16, -2.5673907444456745e-16, -1.734723475976807e-16, -9.020562075079397e-17, 4.163336342344337e-17, -2.498001805406602e-16, -1.1796119636642288e-16, -2.498001805406602e-16, -2.7755575615628914e-17, -1.5959455978986625e-16, -1.457167719820518e-16, -4.163336342344337e-17, -2.3592239273284576e-16, -2.0816681711721685e-16, -8.326672684688674e-17, -1.5959455978986625e-16, -1.942890293094024e-16, -1.734723475976807e-16, -1.8735013540549517e-16, 6.938893903907228e-18, -1.3877787807814457e-17, -3.608224830031759e-16, -1.8735013540549517e-16, -1.3183898417423734e-16, -2.498001805406602e-16, -1.3183898417423734e-16, -2.914335439641036e-16, -1.457167719820518e-16, -3.0531133177191805e-16, -9.020562075079397e-17, -3.0531133177191805e-16, -1.1102230246251565e-16, -1.5959455978986625e-16, -6.245004513516506e-17, -7.632783294297951e-17, -5.551115123125783e-17, -2.0122792321330962e-16, -1.942890293094024e-16, -7.632783294297951e-17, -1.942890293094024e-16, -6.245004513516506e-17, -1.1102230246251565e-16, -1.249000902703301e-16, -1.942890293094024e-16, -1.8735013540549517e-16, -1.5265566588595902e-16, -2.5673907444456745e-16, -1.0408340855860843e-16, -1.0408340855860843e-16, -1.5265566588595902e-16, -1.3183898417423734e-16, -9.71445146547012e-17, -2.2898349882893854e-16, -5.551115123125783e-17, -3.469446951953614e-17]

[glwagner]

The names and comments need to be changed to reflect the notions that:

• grids can be static. Where time-depednent metrics enter, we need a comment / explanation for what happens when the grid is not changing in time
• The data structures for the grid are likely general for any time-dependent vertical coordinate
• The "z*" approximation is a physical / modeling approximation that is independent of the data structures. Specifically, this is a rule for updating the time-dependent grid, plus an approximation that neglects some slope terms, assuming the free surface displacement is small compared to the depth of the domain. However, the "z*" prescription and data structures are properly viewed as independent. This is important for understanding how the code works, plus also for grasping how one would implement other vertical coordinates.