Advecting tracers with parameterized velocities (Mixed-layer Eddies Parameterization)

[iuryt]

Maybe this is going to be solved by #2389 ?

I am trying to implement the mixed-layer eddies parameterization (10.1175/2007JPO3792.1), which

I could make it work to compute Ψeddy at each time step using Callbacks.
https://github.com/iuryt/MLE.jl

resulting in

import xarray as xr
import matplotlib.pyplot as plt

ds = xr.open_dataset("data/output.nc").squeeze()
ds = ds.assign_coords(yC=ds.yC*1e-3)

ti = 0
fig,ax = plt.subplots(2, 1, figsize=(8,8))

fig.subplots_adjust(wspace=0.5)

ds["∂b∂y"].isel(time=ti).plot.contourf(ax=ax[0], cmap="RdBu_r", vmin=-1e-7, vmax=1e-7, levels=12)
ds["Ψx"].isel(time=ti).plot.contourf(ax=ax[1], cmap="PiYG_r", vmin=-12, vmax= 12, levels=11)

for a in ax:
    (-ds["h"]).isel(time=ti).plot(ax=a, color="red")
    ds["b"].isel(time=ti).plot.contour(ax=a, levels=20, colors="k", linestyles="-")
    a.set(
        ylim=[-250,0],
        xlim=[-200,200],
        xlabel="yC [km]",
        title="",
    )

But now I need to also advect tracers with the "extra" velocity field derived from Ψeddy.

What I tried to do

Use BackgroundFields:

So I create some zero background fields for velocities:

ub = Field((Face, Center, Center), grid)
vb = Field((Center, Face, Center), grid)
wb = Field((Center, Center, Face), grid)

model = NonhydrostaticModel(grid = grid,
                            advection = WENO5(),
                            timestepper = :RungeKutta3,
                            coriolis = coriolis,
                            closure = (horizontal_closure, vertical_closure),
                            tracers = (:b),
                            background_fields = (u=ub, v=vb, w=wb),
                            buoyancy = BuoyancyTracer())

Then after creating the simulation, I compute and update background velocities using this:

v_op = @at((Center, Face, Center),   ∂z(Ψx));
w_op = @at((Center, Center, Face), - ∂y(Ψx));

function compute_background_vel!(sim)
    u, v, w = sim.model.background_fields.velocities
    v .= v_op
    w .= w_op
    fill_halo_regions!(v, sim.model.architecture)
    fill_halo_regions!(w, sim.model.architecture)
    return nothing
end
simulation.callbacks[:compute_background_vel] = Callback(compute_background_vel!)

But with the first timestep it returns NaN for resolved u.

Click to see the full code!

using Oceananigans
using Oceananigans.Units

# ---- define size and create a grid

const sponge = 20 #number of points for sponge
const Ny = 46 # number of points in y
const Nz = 48 # number of points in z
const H = 1000 # maximum depth


grid = RectilinearGrid(GPU(),
    size=(Ny+2sponge,Nz),
    halo=(3,3),
    y=(-10*(Ny/2 + sponge)kilometers, 10*(Ny/2 + sponge)kilometers), 
    z=(H * cos.(LinRange(π/2,0,Nz+1)) .- H)meters,
    topology=(Flat, Bounded, Bounded)
)

# ---- 
# ---- rotation

coriolis = FPlane(latitude=60)

# ---- 
# ---- turbulent diffusivity (with sponges)

@inline νh(x,y,z,t) = ifelse((y>-(Ny/2)kilometers)&(y<(Ny/2)kilometers), 1, 100)
horizontal_closure = HorizontalScalarDiffusivity(ν=νh, κ=νh)

@inline νz(x,y,z,t) = ifelse((y>-(Ny/2)kilometers)&(y<(Ny/2)kilometers), 1e-5, 1e-3)
vertical_closure = ScalarDiffusivity(ν=νz, κ=νz)

# ---- 
# ---- instantiate model
ub = Field((Face, Center, Center), grid)
vb = Field((Center, Face, Center), grid)
wb = Field((Center, Center, Face), grid)

model = NonhydrostaticModel(grid = grid,
                            advection = WENO5(),
                            timestepper = :RungeKutta3,
                            coriolis = coriolis,
                            closure = (horizontal_closure, vertical_closure),
                            tracers = (:b),
                            background_fields = (u=ub, v=vb, w=wb),
                            buoyancy = BuoyancyTracer())

# ---- 
# ---- initial conditions

const g = 9.82
const ρₒ = 1026

# background density profile based on Argo data
@inline bg(z) = 0.25*tanh(0.0027*(-653.3-z))-6.8*z/1e5+1027.56

# decay function for fronts
@inline decay(z) = (tanh((z+500)/300)+1)/2

# front function
@inline front(x, y, z, cy) = tanh((y-cy)/12kilometers)

@inline function B(x, y, z)
    fronts = front(x, y, z, -120kilometers) + front(x, y, z, 0) + front(x, y, z, 120kilometers)
    return -( g / ρₒ ) * ( bg(z) + 0.8 * decay(z) * (fronts - 3) / 6 )
end

set!(model; b = B)

# ---- 
# ---- geostrophic initial velocity

b = model.tracers.b
f = model.coriolis.f

# shear operations
uz_op = @at((Face, Center, Center), - ∂y(b) / f );

# compute shear
uz = compute!(Field(uz_op))

# include function for cumulative integration
include("src/cumulative_vertical_integration.jl")

# compute geostrophic velocities
U = cumulative_vertical_integration!(uz)

# prescribe geostrophic velocities for initial condition
set!(model; u = U)

# ---- 
# ---- create a simulation with adaptive time stepping based on CFL condition

simulation = Simulation(model, Δt = 1minutes, stop_time = 1day)

wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=5minutes)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

# ---- 
# ---- compute mixed-layer depth
h = Field{Center, Center, Nothing}(grid)

# buoyancy decrease criterium for determining the mixed-layer depth
const g = 9.82 # gravity
const ρₒ = 1026 # reference density
const Δb=(g/ρₒ) * 0.03

include("src/compute_mixed_layer_depth.jl")
compute_mixed_layer_depth!(simulation) = compute_mixed_layer_depth!(h, simulation.model.tracers.b, Δb)

# add the function to the callbacks of the simulation
simulation.callbacks[:compute_mld] = Callback(compute_mixed_layer_depth!)

# ---- 
# ---- compute horizontal gradient of buoyancy
using Oceananigans.BoundaryConditions: fill_halo_regions!

∂b∂y = Field{Center, Center, Center}(grid)

b = model.tracers.b
∂b∂y_op = ∂y(b);

function compute_∂b∂y!(sim)
    ∂b∂y .= ∂b∂y_op
    fill_halo_regions!(∂b∂y, sim.model.architecture)
    return nothing
end
simulation.callbacks[:compute_∂b∂y] = Callback(compute_∂b∂y!)

# ----
# ---- compute Ψₑ

using KernelAbstractions: @index, @kernel
using KernelAbstractions.Extras.LoopInfo: @unroll
using Oceananigans.Architectures: device_event, architecture
using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Operators: Δzᶜᶜᶜ


@kernel function _compute_Ψₑ!(Ψₑ, grid, h, ∂b, μ, f, ce, Lfₘ, dS)
    i, j = @index(Global, NTuple)

    # average ∇ₕb over the mixed layer
    
    ∂b_sum = 0
    Δz_sum = 0

    h_ij = @inbounds h[i, j]
    
    @unroll for k in grid.Nz : -1 : 1 # scroll from surface to bottom       
        z_center = znode(Center(), Center(), Center(), i, j, k, grid)

        if z_center > -h_ij

            Δz_ijk = Δzᶜᶜᶜ(i, j, k, grid)

            ∂b_sum = ∂b_sum + @inbounds ∂b[i, j, k] * Δz_ijk 
            Δz_sum = Δz_sum + Δz_ijk

        end
    end

    ∂bₘₗ = ∂b_sum/Δz_sum
    
    # compute eddy stream function
    @unroll for k in grid.Nz : -1 : 1 # scroll to point just above the bottom       
        z_center = znode(Center(), Center(), Center(), i, j, k, grid)

        if z_center > -h_ij
            @inbounds Ψₑ[i, j, k] = ce * (dS/Lfₘ) * ((h_ij^2)/abs(f)) * μ(z_center,h_ij) * ∂bₘₗ  
        end
    end

end

function compute_Ψₑ!(Ψₑ, h, ∂b, μ, f; ce = 0.06, Lfₘ = 500meters, dS=10e3)
    grid = h.grid
    arch = architecture(grid)


    event = launch!(arch, grid, :xy,
                    _compute_Ψₑ!, Ψₑ, grid, h, ∂b, μ, f, ce, Lfₘ, dS,
                    dependencies = device_event(arch))

    wait(device_event(arch), event)

    return nothing
end

# x-component of Ψ
Ψx = Field{Center, Center, Center}(grid)

# structure function
@inline μ(z,h) = (1-(2*z/h + 1)^2)*(1+(5/21)*(2*z/h + 1)^2)

compute_Ψₑ!(simulation) = compute_Ψₑ!(Ψx, h, ∂b∂y, μ, f)
# add the function to the callbacks of the simulation
simulation.callbacks[:compute_Ψₑ] = Callback(compute_Ψₑ!)


v_op = @at((Center, Face, Center),   ∂z(Ψx));
w_op = @at((Center, Center, Face), - ∂y(Ψx));

function compute_background_vel!(sim)
    u, v, w = sim.model.background_fields.velocities
    v .= v_op
    w .= w_op
    fill_halo_regions!(v, sim.model.architecture)
    fill_halo_regions!(w, sim.model.architecture)
    return nothing
end
#simulation.callbacks[:compute_background_vel] = Callback(compute_background_vel!)

# ----
# ---- setting up the model output

outputs = merge(model.velocities, model.tracers, (; h, ∂b∂y, Ψx)) # make a NamedTuple with all outputs

simulation.output_writers[:fields] =
    NetCDFOutputWriter(model, outputs, filepath = "data/output.nc",
                     schedule=TimeInterval(8hours))

# ----
# ---- setting up the model progress messages

using Printf

function print_progress(simulation)
    u, v, w = simulation.model.velocities

    # Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, wall time: %s\n",
                   iteration(simulation),
                   prettytime(time(simulation)),
                   prettytime(simulation.Δt),
                   maximum(abs, u), maximum(abs, v), maximum(abs, w),
                   prettytime(simulation.run_wall_time))

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, TimeInterval(1hour))

# ----
# ---- run the simulation

run!(simulation)

[glwagner]

There's a lot to look through so I can't see I've grasped it all, but I do notice that the "logic of the staggered grid" seems to be violated. The other possible issue is that you'll have to ensure that any velocity fields you create have impenetrable boundary conditions where they are supposed to --- this won't happen by default for Field (for example, you wouldn't want "vorticity" to have impenetrable boundary conditions).

On the first point, a v, w stream function should be located at (Center, Face, Face). This is actually crucial for ∂y(v) + ∂z(w) to be machine precision small. You shouldn't have to specify the location of velocities derived from a stream function; they should end up at the right place "naturally". You should also verify that the velocity field derived from a streamfunction is divergence free. Advecting things by divergent velocity field is a sure way to blow things up.

I also saw

∂b∂y = Field{Center, Center, Center}(grid)

But this is a y-gradient of a tracer and should therefore be located at

∂b∂y = Field{Center, Face, Center}(grid)

So I guess in summary the two things I can think of are that the velocity field is either divergent, or advecting things through the boundary (because it doesn't satisfy ImpenetrableBoundaryCondition().

[iuryt]

Initial condition:

After 10 days:

[iuryt]

Note v is impenetrable at north / south (not top / bottom) (and at least I thought you have to declare the location for auxiliary fields (ie fields that aren't passed through model) so

v_mle = FieldBoundaryConditions((Center, Face, Center), grid, north=OpenBoundaryCondition(), south=OpenBoundaryCondition())

similar re: location for w_mle, though the bcs are correct there.

That's totally true, sorry. I just copied and forgot to change.
I will test this as soon as the PR is there. Or.. I can also install a package from github and specify the branch, right?

[glwagner]

It's in Oceananigans now. I had the syntax wrong, the location comes after grid as in

v_mle = FieldBoundaryConditions(grid, (Center, Face, Center), north=OpenBoundaryCondition(), south=OpenBoundaryCondition())

Here's the docstring for FieldBoundaryConditions:

  FieldBoundaryConditions(grid, location; kwargs...)

  Return boundary conditions for auxiliary fields (fields whose values are derived from a model's prognostic fields) on grid and at location.

  Keyword arguments
  ≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡

  Keyword arguments specify boundary conditions on the 6 possible boundaries:

    •  west, left end point in the x-direction where i=1

    •  east, right end point in the x-direction where i=grid.Nx

    •  south, left end point in the y-direction where j=1

    •  north, right end point in the y-direction where j=grid.Ny

    •  bottom, right end point in the z-direction where k=1

    •  top, right end point in the z-direction where k=grid.Nz

  If a boundary condition is unspecified, the default for auxiliary fields and the topology in the boundary-normal direction is used:

    •  PeriodicBoundaryCondition for Periodic directions

    •  GradientBoundaryCondition(0) for Bounded directions and Centered-located fields

    •  nothing for Bounded directions and Face-located fields

    •  nothing for Flat directions and/or Nothing-located fields)

We should throw an error for incorrect usage, but it seems we may not right now.

[iuryt]

Sorry, I meant the AdvectiveForcing, which could be another solution for this problem (maybe better because this is not really a background velocity?).

[glwagner]

I didn't completely understand the example, but you might be right! Yes, you can run code on branches using the package manager eg

pkg> add Oceananigans#glw/forcing-features

you might catch the issue that's prevented merger...

[iuryt]

Ok,

I have advanced a lot with this implementation and I have a test case for a 2d (y-z) front, but this code should work for 3D fronts as well.
My problem now is that I can't figure out why the CPU simulation runs well, but the GPU simulation returns NaN for velocity after 2 days.

I did my best to reduce the code, but I still have to keep kernels for computing mixed-layer depth, ∇b and Ψ.
I am currently using AdvectiveForcing to use the parameterized vecocities to advect buoyancy.
I couldn't figure out what is happening and why it returns NaN only for GPU, only after 2 days of simulation.

Does anyone has a clue on what am I missing here? Sorry for not having a smaller piece of code that reproduces the error.
The code below returns NaN after 2 days of simulation, but runs well if we change the architecture to CPU.

using Oceananigans
using Oceananigans.Units
using Oceananigans.BoundaryConditions: ImpenetrableBoundaryCondition, fill_halo_regions!

# ---- define size and create a grid
const Ny = 46 # number of points in y
const Nz = 48 # number of points in z
const H = 1000 # maximum depth


grid = RectilinearGrid(GPU(),
    size=(Ny, Nz),
    halo=(3, 3),
    y=(-10*(Ny/2)kilometers, 10*(Ny/2)kilometers), 
    z=(H * cos.(LinRange(π/2,0,Nz+1)) .- H)meters,
    topology=(Flat, Bounded, Bounded)
)

# ---- 
# ---- rotation

coriolis = FPlane(latitude=60)

# ---- 
# ---- turbulent diffusivity

horizontal_closure = HorizontalScalarDiffusivity(ν=1, κ=1)
vertical_closure = ScalarDiffusivity(ν=1e-5, κ=1e-5)


# ---- 
# ---- initial mixed-layer eddy velocities

# no penetration bc for mixed-layer eddy vertical velocity
no_penetration = ImpenetrableBoundaryCondition()
w_mle_bc = FieldBoundaryConditions(grid, (Center, Center, Face), top=no_penetration, bottom=no_penetration)

# initial zero fields for mixed-layer eddy velocities
u_mle = XFaceField(grid)
v_mle = YFaceField(grid)
w_mle = ZFaceField(grid, boundary_conditions=w_mle_bc)

# build AdvectiveForcing from mixed-layer eddy velocities
mle_forcing = AdvectiveForcing(u = u_mle, v = v_mle, w = w_mle)

# ---- 
# ---- instantiate model

model = NonhydrostaticModel(grid = grid,
                            advection = WENO5(),
                            timestepper = :RungeKutta3,
                            coriolis = coriolis,
                            closure = (horizontal_closure, vertical_closure),
                            tracers = (:b),
                            forcing = (; b = mle_forcing),
                            buoyancy = BuoyancyTracer())

# ---- 
# ---- initial conditions

const g = 9.82
const ρₒ = 1026

# background density profile based on Argo data
@inline bg(z) = 0.25*tanh(0.0027*(-653.3-z))-6.8*z/1e5+1027.56

# decay function for fronts
@inline decay(z) = (tanh((z+500)/300)+1)/2

# front function
@inline front(x, y, z, cy) = tanh((y-cy)/12kilometers)

@inline function B(x, y, z)
    fronts = front(x, y, z, -120kilometers) + front(x, y, z, 0) + front(x, y, z, 120kilometers)
    return -( g / ρₒ ) * ( bg(z) + 0.8 * decay(z) * (fronts - 3) / 6 )
end

# we start with the fronts but no velocity
# let's leave fronts to adjust
set!(model; b = B)

# ---- 
# ---- create a simulation with adaptive time stepping based on CFL condition

simulation = Simulation(model, Δt = 5minutes, stop_time = 10days)

# ---- 
# ---- compute mixed-layer depth

using KernelAbstractions: @index, @kernel
using KernelAbstractions.Extras.LoopInfo: @unroll
using Oceananigans.Architectures: device_event, architecture
using Oceananigans.Utils: launch!
using Oceananigans.Grids

@kernel function _compute_mixed_layer_depth!(h, grid, b, Δb)
    i, j = @index(Global, NTuple)

    # Use this to set mld to surface
    z_surface = znode(Center(), Center(), Face(), i, j, grid.Nz, grid)
    b_surface = @inbounds b[i, j, grid.Nz] # buoyancy at surface
    
    # at the first iteration, mld is at the surface
    mld_ij = z_surface
    searching_mld = true
    
    @unroll for k in grid.Nz -1 : -1 : 1 # scroll from point just below surface
        # buoyancy and buoyancy decrease
        b_ijk = @inbounds b[i, j, k]
        Δb_ijk = b_surface-b_ijk
        
        # if below mixed layer and we havent yet found it
        if (Δb_ijk>Δb)&(searching_mld)
            # buoyancy and buoyancy decrease for the cell above
            b_ijk_above = @inbounds b[i, j, k-1]
            Δb_ijk_above = b_surface-b_ijk_above
            
            # height for both cells
            z_face = znode(Center(), Center(), Face(), i, j, k, grid)
            z_face_above = znode(Center(), Center(), Face(), i, j, k-1, grid)
            
            # get Δz and Δb
            Δz_ijk = z_face_above-z_face
            
            # interpolate linearly to obtain the mld
            # Note "-" since `h` is supposed to be "depth" rather than "height"
            mld_ij = max(-((Δz_ijk/(Δb_ijk_above-Δb_ijk))*(Δb-Δb_ijk)+z_face), 0)
            
            # stop searching (skip this calculation for other levels)
            searching_mld=false
        end
    end

    @inbounds h[i, j, 1] = mld_ij
end

function compute_mixed_layer_depth!(h, b, Δb)
    grid = h.grid
    arch = architecture(grid)

    event = launch!(arch, grid, :xy,
                    _compute_mixed_layer_depth!, h, grid, b, Δb,
                    dependencies = device_event(arch))

    wait(device_event(arch), event)

    return nothing
end

const g = 9.82 # gravity
const ρₒ = 1026 # reference density

# buoyancy decrease criterium for determining the mixed-layer depth
const Δb=(g/ρₒ) * 0.03

h = Field{Center, Center, Nothing}(grid)
compute_mixed_layer_depth!(simulation) = compute_mixed_layer_depth!(h, simulation.model.tracers.b, Δb)

# add the function to the callbacks of the simulation
simulation.callbacks[:compute_mld] = Callback(compute_mixed_layer_depth!)

# ---- 
# ---- compute gradient of buoyancy

∂b∂x = Field{Face, Center, Center}(grid)
∂b∂y = Field{Center, Face, Center}(grid)
∂b∂z = Field{Center, Center, Face}(grid)

b = model.tracers.b
∂b∂x_op = ∂x(b);
∂b∂y_op = ∂y(b);
∂b∂z_op = ∂z(b);

function compute_∇b!(sim)
    ∂b∂x .= ∂b∂x_op
    ∂b∂y .= ∂b∂y_op
    ∂b∂z .= ∂b∂z_op
    fill_halo_regions!(∂b∂x, sim.model.architecture)
    fill_halo_regions!(∂b∂y, sim.model.architecture)
    fill_halo_regions!(∂b∂z, sim.model.architecture)
    return nothing
end
simulation.callbacks[:compute_∇b] = Callback(compute_∇b!)

# ----
# ---- compute Ψₑ

# repeating the imports just to make it easier while moduling it
using KernelAbstractions: @index, @kernel
using KernelAbstractions.Extras.LoopInfo: @unroll
using Oceananigans.Architectures: device_event, architecture
using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Operators: Δzᶜᶜᶜ, Δzᶜᶜᶠ


@kernel function _compute_Ψₑ!(Ψₑ, grid, h, ∂ₕb, N, μ, f, ce, Lfₘ, ΔS, τ, minpoints, Vm)
    i, j = @index(Global, NTuple)

    # average ∇ₕb and N over the mixed layer
    
    ∂ₕb_sum = 0
    N_sum = 0
    Δz_sum = 0

    h_ij = @inbounds h[i, j]
    
    # number of z points in the mixed layer
    npoints = 0
    
    @unroll for k in grid.Nz : -1 : 1 # scroll from surface to bottom       
        z_center = znode(Center(), Face(), Face(), i, j, k, grid)

        if z_center > -h_ij
            npoints += 1
            
            Δz_ijk = Δzᶜᶜᶜ(i, j, k, grid)

            ∂ₕb_sum = ∂ₕb_sum + @inbounds ∂ₕb[i, j, k] * Δz_ijk
            N_sum = N_sum + @inbounds N[i, j, k] * Δz_ijk 
            Δz_sum = Δz_sum + Δz_ijk

        end
    end

    ∂ₕbₘₗ = ∂ₕb_sum/Δz_sum
    Nₘₗ = N_sum/Δz_sum
    
    Lf = max(Nₘₗ*h_ij/abs(f), Lfₘ)
    
    # compute eddy stream function
    @unroll for k in grid.Nz : -1 : 1 # scroll to point just above the bottom       
        z_face = znode(Center(), Face(), Face(), i, j, k, grid)
        
        Ψ_max = @inbounds Δzᶜᶜᶠ(i, j, k, grid) * Vm
        Ψ_ijk = ce * (ΔS/Lf) * ((h_ij^2)/sqrt(f^2 + τ^-2)) * μ(z_face,h_ij) * ∂ₕbₘₗ
        
        if (z_face > -h_ij)&(npoints > minpoints)
            @inbounds Ψₑ[i, j, k] = min(Ψ_max, Ψ_ijk)
        else
            @inbounds Ψₑ[i, j, k] = 0.0
        end
    end

end

function compute_Ψₑ!(Ψₑ, h, ∂ₕb, N, μ, f; ce = 0.06, Lfₘ = 500meters, ΔS=10e3, τ=86400, minpoints=4, Vm=0.5)
    grid = h.grid
    arch = architecture(grid)


    event = launch!(arch, grid, :xy,
                    _compute_Ψₑ!, Ψₑ, grid, h, ∂ₕb, N, μ, f, ce, Lfₘ, ΔS, τ, minpoints, Vm,
                    dependencies = device_event(arch))

    wait(device_event(arch), event)

    return nothing
end

# create field for each component of Ψ
Ψx = Field{Center, Face, Face}(grid)
Ψy = Field{Face, Center, Face}(grid)

# structure function
@inline μ(z,h) = (1-(2*z/h + 1)^2)*(1+(5/21)*(2*z/h + 1)^2)

# create a function for each component
compute_Ψx!(simulation) = compute_Ψₑ!(Ψx, @at((Center, Face, Nothing), h), ∂b∂y, sqrt(∂b∂z), μ, model.coriolis.f)
compute_Ψy!(simulation) = compute_Ψₑ!(Ψy, @at((Face, Center, Nothing), h), ∂b∂x, sqrt(∂b∂z), μ, model.coriolis.f)

# add the function to the callbacks of the simulation
simulation.callbacks[:compute_Ψx] = Callback(compute_Ψx!)
simulation.callbacks[:compute_Ψy] = Callback(compute_Ψy!)


# ----
# ---- compute mixed-layer eddy velocities

u_mle_op = @at((Face, Center, Center),   ∂z(Ψy));
v_mle_op = @at((Center, Face, Center),   ∂z(Ψx));
w_mle_op = @at((Center, Center, Face), -(∂x(Ψy) + ∂y(Ψx)));

function compute_mle_velocity!(sim)
    u_mle .= u_mle_op
    v_mle .= v_mle_op
    w_mle .= w_mle_op
    Oceananigans.BoundaryConditions.fill_halo_regions!(u_mle)
    Oceananigans.BoundaryConditions.fill_halo_regions!(v_mle)
    Oceananigans.BoundaryConditions.fill_halo_regions!(w_mle)
    return nothing
end


simulation.callbacks[:compute_mle_velocity] = Callback(compute_mle_velocity!)

# ----
# ---- setting up the model output

outputs = merge(model.velocities, model.tracers, (; u_mle=u_mle, v_mle=v_mle, w_mle=w_mle, h, 
                                                    hx=@at((Center, Face, Nothing), h), hy=@at((Face, Center, Nothing), h), 
                                                    ∂b∂x, ∂b∂y, Ψx, Ψy)) # make a NamedTuple with all outputs

simulation.output_writers[:fields] =
    NetCDFOutputWriter(model, outputs, filename = "data/output.nc",
                     schedule=TimeInterval(3hours))

# ----
# ---- setting up the model progress messages

using Printf

function print_progress(simulation)
    u, v, w = simulation.model.velocities

    # Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, Ψmax = (%.1e, %.1e), wall time: %s\n",
                   iteration(simulation),
                   prettytime(time(simulation)),
                   prettytime(simulation.Δt),
                   maximum(abs, u), maximum(abs, v), maximum(abs, w),
                   maximum(abs, Ψx), maximum(abs, Ψy),
                   prettytime(simulation.run_wall_time))

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, TimeInterval(1hour))

# ----
# ---- run the simulation

run!(simulation)

[glwagner]

Here's another tidbit. If you're getting a NaN after one timestep, then this isn't a numerical instability problem --- it means you have 0/0 somewhere. I guess that wouldn't typically be a problem with incorrect boundary conditions...

I also saw that you're computing N with sqrt(∂b∂z). This is a bit dangerous, because ∂b∂z can go slightly negative due to numerical errors. We often use a construct like N = sqrt(max(zero(eltype(grid)), ∂b∂z[i, j, k])). That's probably not your NaN though.

I think maybe paying close attention to the places where you have divisions, and outputting intermediates via @show, could help. It's certainly a mystery why you'd get badness on GPU vs CPU. I think normally that'd mine you're unintentionally indexing something out of bounds, ie accessing junk memory (which is arbitrary, and maybe different "on average" on GPU vs CPU).

[iuryt]

YAHOOO!! It works! It was the N indeed! Thank you so much, I am since yesterday trying to make it work.

[glwagner]

Wow! No idea why that'd be different CPU vs GPU but glad it works.

[iuryt]

I suspect that somehow my GPU is less precise?? Maybe ..

[glwagner]

There's more to it cause you should have gotten DomainError:

julia> sqrt(-1)
ERROR: DomainError with -1.0:

so something's fishy...