Tripolar grid in Oceananigans

[simone-silvestri]

Our implementation of the tripolar grid is in https://github.com/CliMA/OrthogonalSphericalShellGrids.jl at the moment.

We have figured out that the implementation of the tripolar grid is very much entangled with Oceananigans because it requires extending a lot of time-stepping specific code.

This leads to delays in development when we have to develop the two packages in tandem.
Therefore, this PR migrates the implementation of the Tripolar grid and its specific methods to Oceananigans.

[navidcy]

nice

we should update

Oceananigans.jl/docs/src/grids.md

Lines 86 to 93 in dc4a327

	1. [`RectilinearGrid`](@ref Oceananigans.Grids.RectilinearGrid) can be fashioned into lines, rectangles and boxes.
	2. [`LatitudeLongitudeGrid`](@ref Oceananigans.Grids.LatitudeLongitudeGrid) represents sectors of thin spherical shells, with cells bounded by lines of constant latitude and longitude.
	3. [`OrthogonalSphericalShellGrid`](@ref Oceananigans.Grids.OrthogonalSphericalShellGrid) represents sectors of thin spherical shells divided with mesh lines that intersect at right angles (thus, orthogonal) but are otherwise arbitrary.
	!!! note "OrthogonalSphericalShellGrids.jl"
	See the auxiliary package [`OrthogonalSphericalShellGrids.jl`](https://github.com/CliMA/OrthogonalSphericalShellGrids.jl)
	for recipes that implement some useful `OrthogonalSphericalShellGrid`, including the
	["tripolar" grid](https://www.sciencedirect.com/science/article/abs/pii/S0021999196901369).

[navidcy]

@simone-silvestri

Which device is meant to be used here?

Oceananigans.jl/src/Grids/tripolar_grid.jl

Line 108 in db49f4d

loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))

WARNING: both Architectures and CUDA export "device"; uses of it in module Grids must be qualified
Unit tests...: Error During Test at /Users/navid/Research/OC11.jl/test/test_tripolar_grid.jl:31
  Got exception outside of a @test
  UndefVarError: `device` not defined
  Stacktrace:
    [1] (TripolarGrid)(arch::CPU, FT::DataType; size::Tuple{Int64, Int64, Int64}, southernmost_latitude::Int64, halo::Tuple{Int64, Int64, Int64}, radius::Float64, z::Tuple{Int64, Int64}, north_poles_latitude::Int64, first_pole_longitude::Int64)
      @ Oceananigans.Grids ~/Research/OC11.jl/src/Grids/tripolar_grid.jl:108

[glwagner]

Why draft?

[navidcy]

Why draft?

I guess because it needs work... Still most test fail. It's not ready for review...

[glwagner]

Huh true. Just thought it would have been a copy/paste job but maybe there's more to it

[navidcy]

It shouldn't be hard to fix... I can give it a go later!

[simone-silvestri]

the main problem is that we use Fields and fill_halo_regions! in the construction of the grid. I ll search a way around it

[navidcy]

This part:

Oceananigans.jl/src/Grids/tripolar_grid.jl

Lines 139 to 148 in 44f62c1

	# Boundary conditions to fill halos of the coordinate and metric terms
	# We need to define them manually because of the convention in the
	# ZipperBoundaryCondition that edge fields need to switch sign (which we definitely do not
	# want for coordinates and metrics)
	default_boundary_conditions = FieldBoundaryConditions(north = ZipperBoundaryCondition(),
	south = nothing, # The south should be `continued`
	west = Oceananigans.PeriodicBoundaryCondition(),
	east = Oceananigans.PeriodicBoundaryCondition(),
	top = nothing,
	bottom = nothing)

requires BoundaryConditions before grids?

[navidcy]

Oh I just saw your comment @simone-silvestri!

[navidcy]

Damn...

Fields and fill_halo_regions! are also used in the cubed sphere grid also, but that's defined later in the MultiRegion module.

[simone-silvestri]

We could move the zipper-filling halo kernels in grids and use them just as kernels.
Then later import them in boundary conditions and use them for the ZipperBoundaryCondition.
It is a bit convoluted though.

[navidcy]

Or we can define the TripolarGrid in a module on its own after Fields/Boundary conditions? It does make sense that it lives in Grids... aaaaargh

[navidcy]

Btw, here we define a haversine

Oceananigans.jl/src/Grids/tripolar_grid.jl

Lines 457 to 467 in 44f62c1

	# Is this the same as in Oceananigans?
	# TODO: check it out
	function haversine(a, b, radius)
	λ₁, φ₁ = a
	λ₂, φ₂ = b
	x₁, y₁, z₁ = lat_lon_to_cartesian(φ₁, λ₁, radius)
	x₂, y₂, z₂ = lat_lon_to_cartesian(φ₂, λ₂, radius)
	return radius * acos(max(-1.0, min((x₁ * x₂ + y₁ * y₂ + z₁ * z₂) / radius^2, 1.0)))
	end

Elsewhere in Oceananigans we use the Distances.haversine; see https://github.com/JuliaStats/Distances.jl/blob/master/src/haversine.jl

[navidcy]

The two are exactly the same; the snippet below demonstrates that:

using Oceananigans.Grids: lat_lon_to_cartesian
using Distances

# Is this the same as in Oceananigans? 
# TODO: check it out
function haversine_simone(a, b, radius)
    λ₁, φ₁ = a
    λ₂, φ₂ = b

    x₁, y₁, z₁ = lat_lon_to_cartesian(φ₁, λ₁, radius)
    x₂, y₂, z₂ = lat_lon_to_cartesian(φ₂, λ₂, radius)

    return radius * acos(max(-1.0, min((x₁ * x₂ + y₁ * y₂ + z₁ * z₂) / radius^2, 1.0)))
end

using Test

for i in 1:100
    (λ1, φ1) = rand()*360, (2rand()-1)*90
    (λ2, φ2) = rand()*360, (2rand()-1)*90

    radius = rand() * 6300e3

    @test haversine_simone((λ1, φ1), (λ2, φ2), radius) ≈ Distances.haversine((λ1, φ1), (λ2, φ2), radius)
end

[navidcy]

@simone, what's the paper mentioned here:

Note: There is probably a better way to do this, in Murray (2016) they give analytical expressions for the metric terms.

[simone-silvestri]

The date was wrong 😅
The paper is this one https://www.sciencedirect.com/science/article/abs/pii/S0021999196901369

[simone-silvestri]

Let's see if it works with the current workaround

[glwagner]

Should we change how the code is organized? Can we put Fields before Grids?

We could also have a module OrthogonalSphericalShellGrids after Fields. Cubed sphere stuff that is single-panel-specific could go there too.

[simone-silvestri]

Hmmm, I think Grids should be before Fields.
I was thinking to implement the rules to fold the north boundary in the tripolar_grid.jl file

Oceananigans.jl/src/Grids/tripolar_grid.jl

Lines 234 to 260 in ce8c49e

	@inline function fold_north_boundary!(c, i, k, ::Center, ::Center, Nx, Ny, Hy, sign)
	i′ = Nx - i + 1
	for j = 1 : Hy
	@inbounds begin
	c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
	end
	end
	return nothing
	end
	@inline function fold_north_boundary!(c, i, k, ::Face, ::Center, Nx, Ny, Hy, sign)
	i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist!
	s = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign
	i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!!
	for j = 1 : Hy
	@inbounds begin
	c[i, Ny + j, k] = s * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
	end
	end
	return nothing
	end
	@inline function fold_north_boundary!(c, i, k, ::Center, ::Face, Nx, Ny, Hy, sign)

and then use it in boundary conditions:

Oceananigans.jl/src/BoundaryConditions/fill_halo_regions_zipper.jl

Lines 47 to 56 in ce8c49e

	@inline function _fill_north_halo!(i, k, grid, c, bc::ZBC, loc, args...)
	Nx, Ny, _ = size(grid)
	@inbounds ℓx = loc[1]
	@inbounds ℓy = loc[2]
	fold_north_boundary!(c, i, k, ℓx, ℓy, Nx, Ny, Hy, bc.condition)
	return nothing
	end

[glwagner]

Hmmm, I think Grids should be before Fields. I was thinking to implement the rules to fold the north boundary in the tripolar_grid.jl file

Oceananigans.jl/src/Grids/tripolar_grid.jl

Lines 234 to 260 in ce8c49e

	@inline function fold_north_boundary!(c, i, k, ::Center, ::Center, Nx, Ny, Hy, sign)
	i′ = Nx - i + 1
	for j = 1 : Hy
	@inbounds begin
	c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
	end
	end
	return nothing
	end
	@inline function fold_north_boundary!(c, i, k, ::Face, ::Center, Nx, Ny, Hy, sign)
	i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist!
	s = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign
	i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!!
	for j = 1 : Hy
	@inbounds begin
	c[i, Ny + j, k] = s * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
	end
	end
	return nothing
	end
	@inline function fold_north_boundary!(c, i, k, ::Center, ::Face, Nx, Ny, Hy, sign)

and then use it in boundary conditions:

Oceananigans.jl/src/BoundaryConditions/fill_halo_regions_zipper.jl

Lines 47 to 56 in ce8c49e

	@inline function _fill_north_halo!(i, k, grid, c, bc::ZBC, loc, args...)
	Nx, Ny, _ = size(grid)
	@inbounds ℓx = loc[1]
	@inbounds ℓy = loc[2]
	fold_north_boundary!(c, i, k, ℓx, ℓy, Nx, Ny, Hy, bc.condition)
	return nothing
	end

what about having an OrthogonalSphericalShellGrids file after Fields?

It makes sense to me that for non-trivial grids, we actually want to use grids and fields to build the grid. For example, the cubed sphere grid is a conformal map of a rectilinear grid.

Just a thought to save some work but if you think it will ultimate accelerate progress (for some reason) to put everything in Grids then we should do it.

[simone-silvestri]

Well, I guess if we move all the tripolar grid construction in a OrthogonalSphericalShellGrids in a module below Fields and before AbstractOperations, it is just basically a copy and paste of OrthogonalSphericalShellGrids in Oceananigans which indeed, will be very fast.
We could also add the cubed sphere panel there.

I am ok with this solution if you guys think we can do this at this moment. This would simplify all the construction of non-trivial grids in the future (probably the tripolar grid will not be the last).

We could reassess this implementation in the future maybe?

[glwagner]

Well, I guess if we move all the tripolar grid construction in a OrthogonalSphericalShellGrids in a module below Fields and before AbstractOperations, it is just basically a copy and paste of OrthogonalSphericalShellGrids in Oceananigans which indeed, will be very fast. We could also add the cubed sphere panel there.

I am ok with this solution if you guys think we can do this at this moment. This would simplify all the construction of non-trivial grids in the future (probably the tripolar grid will not be the last).

We could reassess this implementation in the future maybe?

I think that's a nice plan for sure. It's nice to have the basic utilities available when building new grids via conformal maps, etc. We can also implement the rotated lat-lon grid there?

[navidcy]

I think it'll be cleaner if indeed we have OrthogonalSphericalShellGrids as a module after Fields (and BoundaryConditions? -- seems like the tripolar needs that as well?)