Pressure Solver Issues with a Multi-Dimensionally-Stretched Grid

[s-aucoin]

Hello,
I am trying to run a simulation with a stretched grid in 3 dimensions. The following code, which has a grid that is only stretched in one dimension, runs without error:

using Oceananigans
using Oceananigans.Units: second, seconds, minute, minutes, hour, hours
using Distributions
using Printf


###############################################################################
## UTILITY FUNCTIONS ##
@inline plateau(z, z₁, z₂) = z >= z₁ && z <= z₂ ? 1.0 : 0.0                          # a plateau function. 1 in the range [z₁, z₂], 0 elsewhere


"""
    HyperbolicCellPositions(k, zₘ, z₀, N, shape_parameter, zc)

Calculate the `z` position as a function of grid cell number `k` using a hyperbolic tangent curve.
The curve is specified by maximum `z` `zₘ`, minimum `z` `z₀`, number of cells `N`, position of minimum grid spacing `zc`, and a `shape parameter`.
"""
function HyperbolicCellPositions(k, zₘ, z₀, N, shape_parameter, zc)

    A = (1 - N)/(tanh(shape_parameter*(z₀ - zc)) - tanh(shape_parameter*(zₘ - zc)))

    D = N - A*tanh(shape_parameter*(zₘ - zc))

    ## Construct the function ##
    return 1/shape_parameter * atanh(1/A * (k - D)) + zc
end



###############################################################################
### UNIVERSAL PARAMETERS ###

# Closure #
const C_SL = 0.16 # Smagorinsky-Lilly Coefficient, use default for now
const ν = 1.0e-6  # kinematic viscosity
const κₛ = 1e-9    # molecular diffusivity of salinity

# Time Parameters #
const Δt_init = 0.01 # initial time step
const Δt_max = 0.05second # maximum time step
const dΔt_max = 1.1 # maximum relative time step change
const max_t = 3minutes # maximum simulation time


###############################################################################
### Define the grid size and dimensions ###
            const Nz = 256  # number of points in the vertical direction
            const Nx = 256  # number of points in the x-direction
            const Ny = 256  # number of points in the y-direction

            ## H/R < 0.8 => R > 1.5 m ##
            const Lz = 1.2  # (m) domain depth above 0
            const H = 0.25  # (m) domain depth below 0
            const Lx = 2.0  # (m) domain length in x
            const Ly = 2.0  # (m) domain length in y

### Stretched grid cell positions ###
shape_p_z = 2.3 # vertical shape parameter
z_cells(k) = HyperbolicCellPositions(k, Lz, -H, Nz, shape_p_z, 0.1)


## Initialize the grid ##
          grid = RectilinearGrid(GPU();
      topology = (Periodic, Periodic, Bounded),
          size = (Nx, Ny, Nz),
             x = (-Lx/2, Lx/2),
             y = (-Ly/2, Ly/2),
             z = z_cells)



###############################################################################
### INITIAL CONDITIONS ###
const σ_b = 2 * LinearEquationOfState().haline_contraction * 9.81                # (m s⁻²) standard deviation for the buoyancy noise. equivalent to 2 psu deficit

@inline Ξ(z, pdf, z₁, z₂) = rand(pdf) * plateau(z, z₁, z₂)                       # Random normally-distributed noise between z₁ and z₂

b_noise_pdf = Normal(0, σ_b)                                                     # set the pdf to sample the random numbers from

bᵢ(x, y, z) = Ξ(z, b_noise_pdf, 0, Lz)                                           # (m s⁻²) The initial buoyancy profile with normalized noise
PTᵢ(x, y, z) = 0                                                                 # Initial tracer distribution


###############################################################################
#### CLOSURE ####
closure = (SmagorinskyLilly(C = C_SL), ScalarDiffusivity(ν=ν, κ=(b=κₛ, PT=κₛ)))


###############################################################################
### DEFINE THE MODEL ###

model = NonhydrostaticModel(; grid,
                            buoyancy = BuoyancyTracer(),
                        advection = WENO(),
                        timestepper = :RungeKutta3,
                            tracers = (:b, :PT),
                            coriolis = FPlane(f=0),
                            closure = closure)



###############################################################################
### INITIAL CONDITIONS ###
set!(model, u=0, v=0, w=0, b=bᵢ, PT=PTᵢ)



###############################################################################
### SET THE SIMULATION UP ###
simulation = Simulation(model, Δt=Δt_init, stop_time=max_t)
wizard = TimeStepWizard(cfl=1.0, max_change=dΔt_max, max_Δt=Δt_max)



###############################################################################
### OUTPUTS ###

# progress message #
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max(|w|) = %.1e m s⁻¹, wall time: %s\n",
                                iteration(sim), prettytime(sim), prettytime(sim.Δt),
                                maximum(abs, sim.model.velocities.w), 
                                prettytime(sim.run_wall_time))

simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(200)) # Print the progress message



###############################################################################
## RUN ##
run!(simulation)

If the code is modified so that the grid is stretched in all 3 dimensions:

### Stretched grid cell positions ###
shape_p_z = 2.3 # vertical shape parameter
shape_p_xy = 2.75 # horizontal shape parameter

x_cells(i) = HyperbolicCellPositions(i, Lx/2, -Lx/2, Nx, shape_p_xy, 0)
y_cells(j) = HyperbolicCellPositions(j, Ly/2, -Ly/2, Ny, shape_p_xy, 0)
z_cells(k) = HyperbolicCellPositions(k, Lz, -H, Nz, shape_p_z, 0.1)


## Initialize the grid ##
          grid = RectilinearGrid(GPU();
      topology = (Bounded, Bounded, Bounded),
          size = (Nx, Ny, Nz),
             x = x_cells,
             y = y_cells,
             z = z_cells)

or 2 dimensions, the simulation fails to find a valid pressure solver:

ERROR: LoadError: None of the implemented pressure solvers for NonhydrostaticModel are supported on 256×256×256 RectilinearGrid{Float64, Bounded, Bounded, Bounded} on CUDAGPU with 3×3×3 halo.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] nonhydrostatic_pressure_solver(arch::GPU{…}, grid::RectilinearGrid{…})
   @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/3ZIHr/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl:58
 [3] nonhydrostatic_pressure_solver(grid::RectilinearGrid{…})
   @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/3ZIHr/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl:62
 [4] NonhydrostaticModel(; grid::RectilinearGrid{…}, clock::Clock{…}, advection::WENO{…}, buoyancy::BuoyancyTracer, coriolis::FPlane{…}, stokes_drift::Nothing, forcing::@NamedTuple{}, closure::Tuple{…}, boundary_conditions::@NamedTuple{}, tracers::Tuple{…}, timestepper::Symbol, background_fields::@NamedTuple{}, particles::Nothing, biogeochemistry::Nothing, velocities::Nothing, hydrostatic_pressure_anomaly::Oceananigans.Models.NonhydrostaticModels.DefaultHydrostaticPressureAnomaly, nonhydrostatic_pressure::Field{…}, diffusivity_fields::Nothing, pressure_solver::Nothing, auxiliary_fields::@NamedTuple{})
   @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/3ZIHr/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:217
 [5] top-level scope
   @ /pool1/data/saucoin/projects/sgd/data/hiji_plume/simulation/tests/stretched_grid/N256/stretched_grid_mwe.jl:100
 [6] include(fname::String)
   @ Base.MainInclude ./client.jl:494
 [7] top-level scope
   @ REPL[1]:1
in expression starting at /pool1/data/saucoin/projects/sgd/data/hiji_plume/simulation/tests/stretched_grid/N256/stretched_grid_mwe.jl:100
Some type information was truncated. Use `show(err)` to see complete types.

I saw that the conclusion from #2516 was that all the implemented pressure solvers were FFT-based in the horizontal for NonhydrostaticModel so they would not work with a stretched grid. I was wondering if this had changed, but I guess the errors I got suggest not? However, if I specify that the simulation should use ConjugateGradientPoissonSolver(), then the simulation doesn't crash, but it never gets past the initial time step.

Is this an issue with the pressure solver failing silently? Is there something I can do to get the simulation to run either with ConjugateGradientPoissonSolver() or another pressure solver?

[glwagner]

Does using

pressure_solver = ConjugateGradientPoissonSolver(grid)

work?

Note, the extra cost of the CG pressure solver will likely dominate the cost savings associated with grid stretching.

[s-aucoin]

If I use

pressure_solver = ConjugateGradientPoissonSolver(grid)

Then the simulation doesn't throw an error, but it did not complete the initial time step after 2 hours of running. After testing with fewer grid points (64x64x64 and 32x32x32), I can say that the simulation does run, but it takes a long time to complete the first time step (~6 minutes and 3.5 seconds respectively), which implies that it would take an extraordinarily long time to run my original simulation. I guess I was just not patient enough.

Thank you!

[glwagner]

Probably that situation can be improved but I am not sure by how much.

Something to explore might be to use the FFT solver in some kind of naive way as a preconditioner (eg building the FFT solver on some proxy uniform grid, and interpolating / regridding between grids?)

The main point perhaps is that this is new territory requiring model development, and at the current moment efficient simulations require two uniform directions so that a direct Poisson solver is available.

[glwagner]

This PR should make your code work automatically: #4720