Channel

[johncmarshall54]

Now that our LES Julia ocean model is emerging, it would be good to srategise about getting a channel going. Important for many reasons but especially for the eddy param problem and also for looking at axymmetric flows, But it seems to me an easy one: we just treat y (across channel) like z. Solver issue is the same etc, am I correct? Science goal would be to train a closure on the x-averaged PV flux. Thoughts, comments? John

[ali-ramadhan]

Yes I think so. If we just impose Neumann boundary conditions in the y-direction then the y-operators look more like the z-operators and the Poisson solver just treats the y-direction with discrete cosine transforms (assuming Δy is constant). @christophernhill might foresee some complications though.

I think @glwagner is working on building a boundary conditions API/abstraction which might be good to have in place first, then implementing this channel model will be easier. But we can prototype something pretty quickly if we want to look at this problem ASAP.

[glwagner]

That'd be amazing to be able to simply drop in walls!

Something that isn't clear to me is what boundary conditions the solver can handle. Can it really handle arbitrary boundary conditions in all directions? I would have thought that the boundary conditions affect the boundary condition on pressure, and thus affect the Poisson solver.

(Or perhaps we would need a new solver, but it's simply a matter of treating the horizontal direction like the vertical one when implementing a new Poisson solver for no-flux boundary conditions in two directions... ?)

Are solver details addressed in the docs?

Something to discuss on Monday...

[ali-ramadhan]

Yes we have some notes on the Fourier Poisson solver here: https://climate-machine.github.io/Oceananigans.jl/latest/algorithm/#Discrete-Fourier-spectral-method-1

Yes we'd need a new Poisson solver with Neumann boundary conditions in two directions but this would be a minor change. I think we'd just have to use the same wavenumbers in the y-direction as we do for the z-direction [see equations (47) and (48)] and use DCTs in the y-direction instead of FFTs.

I think that should be it but not 100% sure if it's that simple.

Note: It's only this simple if Δy and Δz are constants (well, Δz can be variable #46). But if both Δy and Δz are variable, then the Poisson solve will become much slower and more difficult.

[glwagner]

Is the description correct? I thought the solver only uses the Fast Fourier transform algorithm to compute the solution to the discrete Poisson equation (and not the Fourier transform itself). Does the solution method actually rely on the Fourier transform? It's confusing because I believe the wave numbers in 46--48 are incorrect if a spectral method is being used to solve the Poisson equation.

Or am I missing something?

Might be good to add links to the papers that describe this method.

[glwagner]

In summary, according to this original reference, as long as our grid is uniform in all directions, we can implement both Neumann and Dirichlet boundary conditions in all directions. Thus implementing a channel requires a new Poisson solver, but that is very similar to the one we have already implemented for the rigid lid case.

[johncmarshall54]

Is we separate out the hydrostatic pressure - which we are doing - then y is the same as z. Cosines should do it, so it is easy I think. Let's see what Chris and J-M think. It would be good to go after the mesoscale flux problem again with our new tools.

[johncmarshall54]

Test of Channel model, periodic in x, walls in y. Suggest spin-down of a baroclinic zonal flow in thermal wind balance with a cross-stream temperature gradient. Free slip boundary conditions on side walls, top (and also bottom? perhaps no slip at bottom). Initialize with T_y which is constant in x,y and z. and u,v,w=0 at t=0. Add a little noise and watch the front spin down through baroclinic instability. [Note zonal current will come in to geostrophic balance with T_y very fast and then go unstable.] When this is working we can apply a zonal wind-stress.

[ali-ramadhan]

Can confirm that the following Poisson solver generalizes as expected, so we should be able to get to a channel model pretty easily.

For a channel with walls in the y and z directions, the following Poisson solver with proper operators produces a divergence-free solution.

using Test
using FFTW
using Statistics: mean
using GPUifyLoops: @launch, @loop, @synchronize

using Oceananigans

# Increment and decrement integer a with periodic wrapping. So if n == 10 then
# incmod1(11, n) = 1 and decmod1(0, n) = 10.
@inline incmod1(a, n) = ifelse(a==n, 1, a + 1)
@inline decmod1(a, n) = ifelse(a==1, n, a - 1)

@inline δx_c2f(g::RegularCartesianGrid, f, i, j, k) = @inbounds f[i, j, k] - f[decmod1(i, g.Nx), j, k]

@inline function δy_c2f(g::RegularCartesianGrid, f, i, j, k)
    if j == 1
        return 0
    else
        @inbounds return f[i, j, k] - f[i, j-1, k]
    end
end

@inline function δz_c2f(g::RegularCartesianGrid, f, i, j, k)
    if k == 1
        return 0
    else
        @inbounds return f[i, j, k-1] - f[i, j, k]
    end
end

@inline δx²_c2f2c(g::RegularCartesianGrid, f, i, j, k) = δx_c2f(g, f, incmod1(i, g.Nx), j, k) - δx_c2f(g, f, i, j, k)

@inline function δy²_c2f2c(g::RegularCartesianGrid, f, i, j, k)
    if j == g.Ny
        return -δy_c2f(g, f, i, j, k)
    else
        return δy_c2f(g, f, i, j+1, k) - δy_c2f(g, f, i, j, k)
    end
end

@inline function δz²_c2f2c(g::RegularCartesianGrid, f, i, j, k)
    if k == g.Nz
        return δz_c2f(g, f, i, j, k)
    else
        return δz_c2f(g, f, i, j, k) - δz_c2f(g, f, i, j, k+1)
    end
end

@inline function ∇²_pnn(g::RegularCartesianGrid, f, i, j, k)
    (δx²_c2f2c(g, f, i, j, k) / g.Δx^2) + (δy²_c2f2c(g, f, i, j, k) / g.Δy^2) + (δz²_c2f2c(g, f, i, j, k) / g.Δz^2)
end

function ∇²_pnn!(grid::RegularCartesianGrid, f, ∇²f)
    @loop for k in (1:grid.Nz; blockIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds ∇²f[i, j, k] = ∇²_pnn(grid, f, i, j, k)
            end
        end
    end

    @synchronize
end

function solve_poisson_3d_pnn!(g::RegularCartesianGrid, f::CellField, ϕ::CellField)
    kx² = zeros(g.Nx, 1)
    ky² = zeros(g.Ny, 1)
    kz² = zeros(g.Nz, 1)

    for i in 1:g.Nx; kx²[i] = (2sin((i-1)*π/g.Nx)    / (g.Lx/g.Nx))^2; end
    for j in 1:g.Ny; ky²[j] = (2sin((j-1)*π/(2g.Ny)) / (g.Ly/g.Ny))^2; end
    for k in 1:g.Nz; kz²[k] = (2sin((k-1)*π/(2g.Nz)) / (g.Lz/g.Nz))^2; end

    FFTW.r2r!(f.data, FFTW.REDFT10, [2, 3])
    FFTW.fft!(f.data, 1)

    for k in 1:g.Nz, j in 1:g.Ny, i in 1:g.Nx
        @inbounds ϕ.data[i, j, k] = -f.data[i, j, k] / (kx²[i] + ky²[j] + kz²[k])
    end
    ϕ.data[1, 1, 1] = 0

    FFTW.ifft!(ϕ.data, 1)

    @. ϕ.data = real(ϕ.data) / (2g.Ny*2g.Nz)

    FFTW.r2r!(ϕ.data, FFTW.REDFT01, [2, 3])

    nothing
end

function poisson_pnn_planned_div_free_cpu(ft, Nx, Ny, Nz, planner_flag)
    g = RegularCartesianGrid(ft, (Nx, Ny, Nz), (100, 100, 100))

    RHS = CellField(Complex{ft}, CPU(), g)
    RHS_orig = CellField(Complex{ft}, CPU(), g)
    ϕ = CellField(Complex{ft}, CPU(), g)
    ∇²ϕ = CellField(Complex{ft}, CPU(), g)

    RHS.data .= rand(Nx, Ny, Nz)
    RHS.data .= RHS.data .- mean(RHS.data)

    RHS_orig.data .= copy(RHS.data)

    solve_poisson_3d_pnn!(g, RHS, ϕ)
    ∇²_pnn!(g, ϕ, ∇²ϕ)

    ∇²ϕ.data ≈ RHS_orig.data
end

@test poisson_pnn_planned_div_free_cpu(Float64, 16, 16, 16, FFTW.ESTIMATE)

[glwagner]

When implementing the channel model, I think we should consider refactoring our algorithm so that we

1. loop over interior cells that don't touch the boundary;
2. loop over boundary cells taking into account the particular boundary condition.

Note that right now, the algorithm instead is:

A. loop over all cells, treating boundary conditions with branches in low-level operators B. loop over boundary cells to 'adjust' the boundary condition (if the boundary condition is different from what was assumed in A)

I prefer 1 and 2 to A and B. 1 and 2 is also used by other codes, such as Palm:

The model domain is then separated into three subdomains:

A. grid points in free fluid without adjacent surfaces, where the standard PALM code is executed, B. grid points next to surface that require extra code (e.g., surface parametrization), and C. grid points within obstacles, where the standard PALM code is executed but multiplied by zero.

[ali-ramadhan]

Definitely worth thinking about improving our algorithm. With the halo region implementation we have fill_halo_regions!(::Grid, fields...) which is somewhat a step towards your 2, but maybe we don't want to mix boundary conditions and distributed parallelization too much. What Palm does, especially step C, sounds a lot like what the MITgcm does I believe.

For now though I don't have much time and need to change projects for a little while so I'm going to create a channel-model branch and rewrite the operators and Poisson solver to quickly implement a channel model prototype to play around with a bit. Then I'll revisit the channel-model branch when I have more time to build the appropriate abstractions, e.g.

abstract type GridBoundaryType end
struct Periodic <: GridBoundaryType end
struct Bounded <: GridBoundaryType end
Grid{A<:AbstractArray, R<:AbstractRange, BX<:GridBoundaryType, BY<:GridBoundaryType, BZ<:GridBoundaryType}

and properly merge it into master with tests.

[ali-ramadhan]

@ali-ramadhan, I don't think we should implement boundary conditions via operators. I think this solution is difficult, because the mapping from operators to boundary conditions is not straightforward. It is also potentially inflexible, because in general we want to allow different boundary conditions to be applied to different fields. It isn't clear to me how we would implement different boundary conditions on left and right, for example.

I think there are two good solutions to the boundary conditions issue. One is to implement the boundary conditions with the halo regions. Note that this is currently the strategy we use for the horizontal boundary conditions. The idea is to fill in the halo regions such that applying the RHS kernel to the boundary elements produces the correct boundary condition / fluxes. Since we already fill halo regions, this method does not have any additional cost. This is also the method used in OceanTurb. Another benefit of this method is that the boundary conditions can actually be saved to disk along with the interior part of the solution, and the value of fields on the boundary can be calculated directly from the saved field.

Originally posted by @glwagner in https://github.com/climate-machine/Oceananigans.jl/pull/283#issuecomment-500821917

[ali-ramadhan]

Reason being is that apparently there are issues or reasons why it's better to enforce the wall by setting the area to zero at the wall, rather than using the halo regions to enforce things. I still need to discuss this with @jm-c and @christophernhill.

@ali-ramadhan when we discussed this long ago, the reasoning was that this method did not generalize to the particular case where the rough bathymetry was organized in such a way that a single halo point needed to be set to determine fluxes at two different fluid points.

I think that this concern is specific to a particular method of enforcing bathymetry (ie, the one similar to that implemented in MITgcm).

Another method is to apply the RHS kernel only to interior cells that are not adjacent to a boundary, and then to run an additional loop over the cells that are adjacent to the (possibly rough) boundary to add the appropriate fluxes. This is the method used in PALM.

I feel that the second method (not using halo regions) is more complex. Personally I prefer using the halo regions for now because it seems that bathymetry is relatively far in the future at this point, and its best to focus now at the task at hand. We can reconsider the boundary condition issue when we decided to add bathymetry later.

Originally posted by @glwagner in https://github.com/climate-machine/Oceananigans.jl/pull/283#issuecomment-500824068

[ali-ramadhan]

@glwagner I moved your comments about implementing walls from PR #283 to this issue where they are relevant.

[ali-ramadhan]

I don't have enough experience to say which method is better so I'll ask @jm-c and @christophernhill to learn more and reply back. We're also thinking of adding MPI support sooner rather than later (same with supporting multiple models per GPU) which may or may not change things.

Another benefit of this method is that the boundary conditions can actually be saved to disk along with the interior part of the solution, and the value of fields on the boundary can be calculated directly from the saved field.

Just noting that whether we should save boundary values (i.e. halo data) to disk is an open issue #92.

[glwagner]

Just noting that whether we should save boundary values (i.e. halo data) to disk is an open issue #92.

The user can choose how to save their data in general, so the point holds regardless of how the 'default' output writer is designed.