Correctness test for `update_hydrostatic_pressure!`

[glwagner]

Here's a correctness test for a single kernel (update_hydrostatic_pressure!)

using Oceananigans
using Oceananigans.Architectures: ReactantState
using Reactant
using Statistics

function simple_model(arch)
    Nx = Nz = 16
    grid = RectilinearGrid(arch, size=(Nx, 1, Nz), x=(-Nx/2, Nx/2), y=(0, 1), z=(0, Nz))
    free_surface = ExplicitFreeSurface()
    model = HydrostaticFreeSurfaceModel(; grid, free_surface, buoyancy=BuoyancyTracer(), tracers=:b)
    bᵢ(x, y, z) = exp(-x^2)
    set!(model, b=bᵢ)
    return model
end

c_model = simple_model(CPU())
r_model = simple_model(Oceananigans.Architectures.ReactantState())

using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!

c_p = c_model.pressure.pHY′
r_p = r_model.pressure.pHY′
@time update_hydrostatic_pressure!(c_p, CPU(), c_model.grid, c_model.buoyancy, c_model.tracers)

r_update_hydrostatic_pressure! = @compile(
    sync = true,
    raise = true,
    update_hydrostatic_pressure!(
        r_p,
        ReactantState(),
        r_model.grid,
        r_model.buoyancy,
        r_model.tracers
    )
)

@time r_update_hydrostatic_pressure!(r_p, ReactantState(), r_model.grid, r_model.buoyancy, r_model.tracers)

pc = c_model.pressure.pHY′
pr = r_model.pressure.pHY′

pcp = interior(pc, :, :, 1)
pcr = Array(interior(pr, :, :, 1))

@show maximum(abs, pcp .- pcr)

gives

julia> include("mwe.jl")
  0.001675 seconds (1.46 k allocations: 116.469 KiB, 91.55% compilation time)
  0.011373 seconds (48.76 k allocations: 3.376 MiB, 99.31% compilation time)
maximum(abs, pcp .- pcr) = 0.38940039153570183
0.38940039153570183

My status

(GordonBell25) pkg> st
Project GordonBell25 v0.1.0
Status `~/Projects/GB-25/Project.toml`
  [6e4b80f9] BenchmarkTools v1.6.0
  [179af706] CFTime v0.1.4
  [0376089a] ClimaOcean v0.5.6
  [9e8cae18] Oceananigans v0.96.0 `https://github.com/CliMA/Oceananigans.jl.git#wm/ostracer`
  [aea7be01] PrecompileTools v1.2.1
  [3c362404] Reactant v0.2.47
  [efcf1570] Setfield v1.1.2
⌅ [0234f1f7] HDF5_jll v1.14.3+3
  [0192cb87] Reactant_jll v0.0.95+0
  [ade2ca70] Dates
  [de0858da] Printf
  [9abbd945] Profile
  [9a3f8284] Random
  [9e88b42a] Serialization

This result IS different than with Reactant_jll 0.0.94. I'm just guessing off the top of my head but I think the max discrepancy was ~12.

[glwagner]

Note I ran with CliMA/Oceananigans.jl#4242 but this shouldn't make a difference

[glwagner]

many artifacts from baroclinic_instability_run

generated with this script https://github.com/PRONTOLab/GB-25/blob/7d431149065eda7d97ec71d114c291db66c7977d/ocean-climate-simulation/baroclinic_instability_run.jl

[glwagner]

Here's another test that runs a model:

using Oceananigans
using Oceananigans.Architectures: ReactantState
using Reactant
using Statistics

function first_time_step!(model)
    Δt = model.clock.last_Δt
    Oceananigans.TimeSteppers.first_time_step!(model, Δt)
    return nothing
end

function loop!(model, Ninner)
    Δt = model.clock.last_Δt
    @trace for _ = 1:Ninner
        Oceananigans.TimeSteppers.time_step!(model, Δt)
    end
    return nothing
end

function simple_model(arch, free_surface)
    Ny = Nz = 16
    φ₀ = 40
    grid = LatitudeLongitudeGrid(arch,
                                 topology = (Periodic, Bounded, Bounded),
                                 size = (Ny, Ny, Nz),
                                 longitude = (-10, 10),
                                 latitude = (φ₀ - 10, φ₀ + 10),
                                 z = (-1000, 0),
                                 halo = (6, 6, 6))

    model = HydrostaticFreeSurfaceModel(; grid, free_surface,
                                        buoyancy = BuoyancyTracer(),
                                        tracers = :b)

    model.clock.last_Δt = 1

    bᵢ(x, y, z) = exp(-x^2)
    set!(model, b=bᵢ)

    return model
end

for free_surface = [SplitExplicitFreeSurface(substeps=10), ExplicitFreeSurface()]

    c_model = simple_model(CPU(), free_surface)
    r_model = simple_model(ReactantState(), free_surface)

    @info "Compiling..."
    r_update_state! = @compile sync=true raise=true Oceananigans.TimeSteppers.update_state!(r_model)
    r_first_time_step! = @compile sync=true raise=true first_time_step!(r_model)
    r_loop! = @compile sync=true raise=true loop!(r_model, ConcreteRNumber(10))

    first_time_step!(c_model)
    r_first_time_step!(r_model)

    c_η = c_model.free_surface.η
    r_η = r_model.free_surface.η

    c_u = c_model.velocities.u
    r_u = r_model.velocities.u

    c_p = c_model.pressure.pHY′
    r_p = r_model.pressure.pHY′

    c_up = Array(interior(c_u))
    r_up = Array(interior(r_u))

    c_pp = Array(interior(c_p))
    r_pp = Array(interior(r_p))

    @show free_surface
    @show maximum(abs, c_ηp .- r_ηp)
    @show maximum(abs, c_up .- r_up)
    @show maximum(abs, c_pp .- r_pp)
end

result

[ Info: Compiling...
free_surface = SplitExplicitFreeSurface substepping with FixedSubstepNumber(7)

maximum(abs, c_ηp .- r_ηp) = 0.0
maximum(abs, c_up .- r_up) = 0.003831976957230366
maximum(abs, c_pp .- r_pp) = 1.734723475976807e-18
[ Info: Compiling...
free_surface = ExplicitFreeSurface{Nothing, Float64}(nothing, 9.80665)
maximum(abs, c_ηp .- r_ηp) = 0.0
maximum(abs, c_up .- r_up) = 8.673617379884035e-19
maximum(abs, c_pp .- r_pp) = 1.734723475976807e-18

This test comes out differently than the first MWE. pressure is fine --- inconsistent with the first test --- but u is wrong. Also the problem depends on the type of free_surface, which could hint that SplitExplicitFreeSurface is a culprit. That's implemented here

https://github.com/CliMA/Oceananigans.jl/tree/main/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces

[simone-silvestri]

This works, I ll try to take a look at the problem that fails

using Oceananigans
using Oceananigans.Architectures: ReactantState
using Reactant
using Statistics
using Test

using Oceananigans.Models.HydrostaticFreeSurfaceModels: step_free_surface!

function first_time_step!(model)
    Δt = model.clock.last_Δt
    Oceananigans.TimeSteppers.first_time_step!(model, Δt)
    return nothing
end

function loop!(model)
    Δt = model.clock.last_Δt
    for _ in 1:10
        Oceananigans.TimeSteppers.time_step!(model, Δt)
    end
end

const u_init = rand(16, 16, 16)
const b_init = rand(16, 16, 16) * 0.001

function simple_model(arch, free_surface)
    Ny = Nz = 16
    φ₀ = 40
    grid = LatitudeLongitudeGrid(arch,
                                 topology = (Periodic, Bounded, Bounded),
                                 size = (Ny, Ny, Nz),
                                 longitude = (-10, 10),
                                 latitude = (φ₀ - 10, φ₀ + 10),
                                 z = (-1000, 0),
                                 halo = (6, 6, 6))

    model = HydrostaticFreeSurfaceModel(; grid, free_surface,
                                          buoyancy = BuoyancyTracer(),
                                          tracers = :b)

    model.clock.last_Δt = 1
    set!(model, u=u_init, b=b_init) 

    return model
end

free_surface = SplitExplicitFreeSurface(substeps=10)

c_model = simple_model(CPU(), free_surface)
r_model = simple_model(ReactantState(), free_surface)

@show @test r_model.grid == r_model.free_surface.η.grid

@info "Compiling..."
r_barotropic_step! = @compile sync=true raise=true barotropic_step!(r_model)
r_update_state!    = @compile sync=true raise=true Oceananigans.TimeSteppers.update_state!(r_model)
r_first_time_step! = @compile sync=true raise=true first_time_step!(r_model)
r_loop!            = @compile sync=true raise=true loop!(r_model)

first_time_step!(c_model)
r_first_time_step!(r_model)
loop!(c_model)
r_loop!(r_model)

c_η = c_model.free_surface.η
r_η = r_model.free_surface.η

c_u = c_model.velocities.u
r_u = r_model.velocities.u

c_p = c_model.pressure.pHY′
r_p = r_model.pressure.pHY′

c_ηp = Array(interior(c_η))
r_ηp = Array(interior(r_η))

c_up = Array(interior(c_u))
r_up = Array(interior(r_u))

c_pp = Array(interior(c_p))
r_pp = Array(interior(r_p))

@show free_surface
@show maximum(abs, c_ηp .- r_ηp)
@show maximum(abs, c_up .- r_up)
@show maximum(abs, c_pp .- r_pp)