How to create and save field-derived fields (e.g Richardson Number) and how to use it for a Turbulence Closure

[iuryt]

How to create field derived fields? For instance, if I want the model to calculate and save Richardson number fields.

Also, I saw that the diffusivity can be a function of x,y,z,t, can I easily give a function that depend on the Richardson number or do I need to implement it as a Turbulence closure?

[glwagner]

@iuryt this is a great question! Do you mind if I convert this to a discussion? I think there are multiple answers and a discussion might be better suited for keeping track of the solutions for future users.

There are a few possible solutions. Note that the diffusivities can also be AbstractArray, which include Field and AbstractOperations. Using either a concrete Field or AbstractOperation is how you'll solve this problem.

Solution: pre-define a Richardson number field and compute in Simulation.callbacks

This is perhaps the simplest solution: create a field Ri = Field{Center, Center, Face}(grid), and then to define an abstract operation that's a function of this field as your diffusivity, something like:

# Pacanowski-Philander (eg https://glwagner.github.io/OceanTurb.jl/latest/models/pacanowskiphilander/)
Ri = Field{Center, Center, Face}(grid) # ∂z(b) / (∂z(u)^2 + ∂z(v)^2)
ν₀ = 1e-4
ν₁ = 1e-2
κ₀ = 1e-5
κ₁ = 1e-2
c = 5
n = 2

ν = @at (Center, Center, Center) ν₀ + ν₁ / (1 + c * Ri^n)
κᵇ = @at (Center, Center, Center) κ₀ + κ₁ / (1 + c * Ri^(n+1))

@show ν # it's an operation...

closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(); ν, κ=(; b=κᵇ))

Then, you need to compute Ri in a Simulation.callback (you also need to fill it's halo regions, which is a bit annoying but we show you how below). This is pretty advanced usage (and there are subtleties as I mention below) so I'll write a script that implements Pacanowski-Philander as an example. Note that it is currently a limitation of our turbulence closures that viscosities diffusivities must be defined at (Center, Center, Center). #2295 makes some changes that will permit this restriction to be relaxed in the future, but we aren't quite there yet.

Solution 2: precreate velocities and tracers

Another solution is to pre-create velocities and tracers by pre-creating a "fake model". Then, the viscosity / diffusivities can be defined directly as abstract operations of fake_model.velocities and fake_model.tracers, and these named tuples can be passed on to a "true" model. This could be slightly more efficient than computing and storing Ri every time-step. This would look something like

grid = RectilinearGrid(size=128, z=(-100, 0), halo=3, topology=(Flat, Flat, Bounded))
fake_model = NonhydrostaticModel(; grid, tracers=:b, buoyancy=nothing)

velocities = fake_model.velocities
tracers = fake_model.tracers
Ri = ∂z(b) / (∂z(u)^2 + ∂z(v)^2)

ν₀ = 1e-4
ν₁ = 1e-2
κ₀ = 1e-5
κ₁ = 1e-2
c = 5
n = 2

ν = @at (Center, Center, Center) ν₀ + ν₁ / (1 + c * Ri^n)
κᵇ = @at (Center, Center, Center) κ₀ + κ₁ / (1 + c * Ri^(n+1))
closure = ScalarDiffusivity(; ν, κ=(; b=κᵇ))
true_model = NonhydrostaticModel(; grid, closure, velocities, tracers, buoyancy=BuoyancyTracer())

which produces

julia> true_model = NonhydrostaticModel(; grid, closure, velocities, tracers, buoyancy=BuoyancyTracer())
NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
├── grid: 1×1×128 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
├── timestepper: QuasiAdamsBashforth2TimeStepper
├── tracers: b
├── closure: ScalarDiffusivity{ExplicitTimeDiscretization}(ν=BinaryOperation at (Center, Center, Center), κ=(b=UnaryOperation at (Center, Center, Center),))
├── buoyancy: BuoyancyTracer with -ĝ = ZDirection
└── coriolis: Nothing

note the type of model.closure. (Also note this doesn't work with HydrostaticFreeSurfaceModel apparently... we need to fix that.)

One downside of this method is that it precludes VerticallyImplicitTimeDiscretization(): because we use a predictor-corrector method for implicit time discretization, diffusivities that are treated implicitly must be precomputed.

Subtleties

There's more. There are subtleties peculiar to using Ri in a closure, because Ri easily becomes NaN when both the buoyancy gradient and shear are zero. The usual solution to this for ocean modeling is to use the "hack" that 0/0 = 0. In other words, NaNs are intercepted and set to 0. There are again a few solutions to this. One is simply to set all NaN to zero in side the function that computes Ri, eg

u, v, w = model.velocities
b = model.tracers.b
Ri_op = @at (Center, Center, Center) ∂z(b) / (∂z(u)^2 + ∂z(v)^2)

""" Compute the Richardson number and store in `Ri`. """
function compute_Ri!(sim)
    Ri .= Ri_op

    # Zero out NaNs
    Ri_parent = parent(Ri)
    parent(Ri_parent)[isnan.(Ri_parent)] .= 0

    fill_halo_regions!(Ri, sim.model.architecture)
    return nothing
end

Other possibilities are 1) use ConditionalOperation to intercept 0/0 and set them to 0 during the Ri computation (I don't have the code for that off the top of my head so I'll have to look into it and post later) or 2) write a kernel function (this requires using more primitive operators than AbstractOperations, and also being careful about staggered grid locations) to use with KernelFunctionOperation that does the necessary checking. There's a lot of great examples of how to use KernelFunctionOperation here: https://github.com/tomchor/Oceanostics.jl.

[glwagner]

If there are ways to make it easier (eg what would your "dream" code look like?) then we should pursue them. This is exactly the kind of thing that I think we want to make as easy as possible (and we can make as easy as possible) with Oceananigans.

[glwagner]

This script:

using Oceananigans
using Oceananigans.Units
using Oceananigans.BoundaryConditions: fill_halo_regions!
using GLMakie

grid = RectilinearGrid(size=128, z=(-100, 0), halo=3, topology=(Flat, Flat, Bounded))

# Pacanowski-Philander (eg https://glwagner.github.io/OceanTurb.jl/latest/models/pacanowskiphilander/)
Ri = Field{Center, Center, Center}(grid) # ∂z(b) / (∂z(u)^2 + ∂z(v)^2)
ν₀ = 1e-4
ν₁ = 1e-2
κ₀ = 1e-5
κ₁ = 1e-2
c = 5
n = 2

ν = ν₀ + ν₁ / (1 + c * Ri^n)
κᵇ = κ₀ + κ₁ / (1 + c * Ri^(n+1))

@show ν # it's an operation...

closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(); ν, κ=(; b=κᵇ))

model = HydrostaticFreeSurfaceModel(; grid, closure, buoyancy=BuoyancyTracer(), tracers=:b)
simulation = Simulation(model, Δt=1minute, stop_time=1day)

# Define an AbstractOperation that computes Ri in advance
u, v, w = model.velocities
b = model.tracers.b
Ri_op = @at (Center, Center, Center) ∂z(b) / (∂z(u)^2 + ∂z(v)^2)

""" Compute the Richardson number and store in `Ri`. """
function compute_Ri!(sim)
    Ri .= Ri_op

    # Zero NaNs
    Ri_parent = parent(Ri)
    parent(Ri_parent)[isnan.(Ri_parent)] .= 0

    fill_halo_regions!(Ri, sim.model.architecture)
    return nothing
end

simulation.callbacks[:compute_Ri] = Callback(compute_Ri!) # uses default IterationInterval(1)

fields = []
function grab_fields!(sim)
    uⁿ = Array(interior(u, 1, 1, :))
    vⁿ = Array(interior(v, 1, 1, :))
    bⁿ = Array(interior(b, 1, 1, :))
    iter = iteration(sim)
    t = time(sim)
    push!(fields, (; t, iter, u=uⁿ, v=vⁿ, b=bⁿ))
    return nothing
end

simulation.callbacks[:grabber] = Callback(grab_fields!, TimeInterval(10minutes))

N² = 1e-4
S² = 1e-3 # Ri = 0.1
z₀ = -50
Δzu = 4
Δzb = 16

step(x, c, w) = 1/2 * (1 + tanh((x - c) / w)) # smooth step function

uᵢ(x, y, z) = Δzu * sqrt(S²) * step(z, z₀, Δzu)
bᵢ(x, y, z) = Δzb * N² * step(z, z₀, Δzb)

set!(model, u=uᵢ, b=bᵢ)

run!(simulation)

z = znodes(Center, grid)
fig = Figure()
ax_u = Axis(fig[1, 1], ylabel="z (m)", xlabel="u (m s⁻¹)")
ax_b = Axis(fig[1, 2], ylabel="z (m)", xlabel="b (m s⁻²)")
slider = Slider(fig[2, :], range=1:length(fields), startvalue=1)
n = slider.value
uⁿ = @lift fields[$n].u
bⁿ = @lift fields[$n].b

lines!(ax_u, uⁿ, z)
lines!(ax_b, bⁿ, z)

title = @lift "Diffusing shear layer at t = " * prettytime(fields[$n].t)
Label(fig[0, :], title)

display(fig)

record(fig, "pacanowski_philander_diffusion.gif", 1:length(fields), framerate=12) do nn
    @info "Drawing frame $nn of $(length(fields))..."
    n[] = nn
end

produces

My main concern is that our broadcasting may not be performant when used online during a simulation. Possibly, we can fix that or make it better.

[glwagner]

Ok, a bit better now!

The following uses some lower-level Oceananigans functions and KernelFunctionOperation. I suspect this is a bit more performant, but I'm not sure. The main reason is that I think our broadcasting machinery has some overhead right now, so writing field .= op is not all that cheap (we can fix this, but it might take some work). @iuryt if you have the chance to benchmark different solutions, it'd be interesting to hear what works best!

The main downside of this solution is that we need to understand the staggered grid to implement it. Maybe not too onerous (@simone-silvestri thinks I should give users more credit), but part of me feels like we should be able to auto-magic our way around this. The main barrier to using abstract operations here is figuring out how to implement this function with a ConditionalOperation (and also having < as a valid BinaryOperation, eg figuring out #2169).

I also like the following because it's the first known example of auxiliary_fields being used. Hooray for that! Also I realized that we can just use a linear stratification which is nice. I made the shear stronger to increase the drama.

Working on this helped uncover a few wrinkles in the user API:

• closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(); ν, κ=κᵇ) doesn't work (Can't use single ScalarDiffusivity(; κ) when κ isa AbstractArray #2342)
• HydrostaticFreeSurfaceModel(; velocities=velocities) doesn't work (Can't use kwarg velocities with HydrostaticFreeSurfaceModel for pre-allocated fields #2341)

using Oceananigans
using Oceananigans.Units
using Oceananigans.Operators
using GLMakie

# A bit of code...
@inline f²(i, j, k, grid, f, args...) = @inbounds f(i, j, k, grid, args...)^2

@inline function Riᶜᶜᶜ(i, j, k, grid, U, b)
    N² = ℑzᵃᵃᶜ(i, j, k, grid, ∂zᶜᶜᶠ, b)
    S²u = ℑxzᶜᵃᶜ(i, j, k, grid, f², ∂zᶠᶜᶠ, U.u)
    S²v = ℑyzᵃᶜᶜ(i, j, k, grid, f², ∂zᶜᶠᶠ, U.v)
    S² = S²u + S²v
    return ifelse(S² == 0, zero(eltype(grid)), N² / S²)
end

grid = RectilinearGrid(size=128, z=(-100, 0), halo=3, topology=(Flat, Flat, Bounded))
fake_model = HydrostaticFreeSurfaceModel(; grid, tracers=:b, buoyancy=BuoyancyTracer())
velocities = fake_model.velocities
tracers = fake_model.tracers

# Pacanowski-Philander implementation
#
# The following implements the Packanowski-Philander model for shear-modulated mixing
# with parameters:
ν₀ = 1e-4
ν₁ = 1e-2
κ₀ = 1e-5
κ₁ = 1e-2
c = 5
n = 2

# In Packanowski-Philander both the viscosity and diffusivity
# depend on the Richardson number:
Ri_op = KernelFunctionOperation{Center, Center, Center}(Riᶜᶜᶜ, grid, computed_dependencies=(velocities, tracers.b))
Ri = Field(Ri_op)

ν  = ν₀ + ν₁ / (1 + c * Ri^n)
κᵇ = κ₀ + κ₁ / (1 + c * Ri^(n+1))

@show ν # it's an operation...

closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(); ν, κ=(; b=κᵇ))

# For `closure` to work correctly, the Ri must be defined as an "auxiliary field" of the model.
# Fields in model.auxiliary_fields are updated every time-stepper stage.
model = HydrostaticFreeSurfaceModel(; grid, closure, tracers=:b, buoyancy=BuoyancyTracer(), auxiliary_fields = (; Ri))
model.velocities = velocities
model.tracers = tracers

# Initial condition with Ri ≈ 0.01
step(x, c, w) = 1/2 * (1 + tanh((x - c) / w)) # smooth step function

N² = 1e-4
bᵢ(x, y, z) = N² * z

S² = 1e-2
Δu = 4  # m
uᵢ(x, y, z) = Δu * sqrt(S²) * step(z, -grid.Lz/2, Δu)

set!(model, u=uᵢ, b=bᵢ)

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

# Alternative to writing output.
fields = []
function grab_fields!(sim)
    u, v, w = sim.model.velocities
    b = sim.model.tracers.b
    Ri = sim.model.auxiliary_fields.Ri
    uⁿ = Array(interior(u, 1, 1, :))
    vⁿ = Array(interior(v, 1, 1, :))
    bⁿ = Array(interior(b, 1, 1, :))
    Riⁿ = Array(interior(Ri, 1, 1, :))
    push!(fields, (; t=time(sim), i=iteration(sim), u=uⁿ, v=vⁿ, b=bⁿ, Ri=Riⁿ))
    return nothing
end

simulation.callbacks[:grabber] = Callback(grab_fields!, TimeInterval(10minutes))

run!(simulation)

# Plot results
z = znodes(Center, grid)
fig = Figure()
ax_u = Axis(fig[1, 1], ylabel="z (m)", xlabel="u (m s⁻¹)")
ax_b = Axis(fig[1, 2], ylabel="z (m)", xlabel="b (m s⁻²)")
ax_R = Axis(fig[1, 3], ylabel="z (m)", xlabel="Richardson number")
slider = Slider(fig[2, :], range=1:length(fields), startvalue=1)
n = slider.value
uⁿ = @lift fields[$n].u
bⁿ = @lift fields[$n].b
Riⁿ = @lift fields[$n].Ri

lines!(ax_u, uⁿ, z)
lines!(ax_b, bⁿ, z)
lines!(ax_R, Riⁿ, z)
xlims!(ax_R, 0, 5.0)

title = @lift "Diffusing shear layer at t = " * prettytime(fields[$n].t)
Label(fig[0, :], title)

display(fig)

record(fig, "pacanowski_philander_diffusion.gif", 1:length(fields), framerate=12) do nn
    @info "Drawing frame $nn of $(length(fields))..."
    n[] = nn
end

We probably put this or something like it into the docs at some point

.

[iuryt]

@glwagner
I am really impressed with your effort to explain to me all details of the implementation. Thanks you!
In the past few days I am learning a lot reading the code/documentation and examples of Oceananigans.

When I started this issue, I wasn't really expecting to have Pacanowski-Philander parameterization implemented in a few hours. Kudos to you!!
I know this parameterization is very simple, but implementing it on a FORTRAN-based GCM from scratch would take waaaay more time, at least for me.

Can you help me to see if I understood everything correctly? Please, sorry if I am asking too-trivial questions.

Why you need to define f², wouldn't work simply using ∂zᶠᶜᶠ and squaring the operation?

I understand that you create a fake model and don't run it, simply use it to update the auxiliary field. I don't understand why you could not do the same for the main model. For instance, you could first create the model, define Ri and closure and then set! the model with this new closure and auxiliary field. Does that makes sense? I understand the in the background this might not work the way it is now. Something like this:

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

Ri = Field(KernelFunctionOperation{Center, Center, Center}(Riᶜᶜᶜ, grid, computed_dependencies=(model.velocities, model.tracers.b)))

ν  = ν₀ + ν₁ / (1 + c * Ri^n)
κᵇ = κ₀ + κ₁ / (1 + c * Ri^(n+1))

closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(); ν, κ=(; b=κᵇ))

set!(model, closure = closure, auxiliary_fields = (; Ri))

Thinking on a more general way, it would be awesome to give a general function that could depend on any field in the model, including auxiliary. I am thinking that this way it would be easier for users to create and implement their own parameterizations. Including using any model in general that could come in a form of a function. For instance, a ML model could be simply used as a turbulence closure.

Again, thanks for all attention you are giving to me. As a new user, it makes me learn more and more and maybe sometime I can also retribute by contributing to the model development.

[glwagner]

When I started this issue, I wasn't really expecting to have Pacanowski-Philander parameterization implemented in a few hours. Kudos to you!!

Thanks for asking the question! Also is it ok if I convert to discussion? If it remains an issue we might lose track of it when the issue is closed. But this might be useful to people if its more visible.

Why you need to define f², wouldn't work simply using ∂zᶠᶜᶠ and squaring the operation?

Because we need to interpolate ∂zᶠᶜᶠ using ℑxzᶜᵃᶜ. Note that Oceananigans operators (derivatives, interpolation, differences) are composable. The function f² is used in this composition --- it squares a function, rather than a field. So ℑxzᶜᵃᶜ(i, j, k, grid, f², ∂zᶠᶜᶠ, U.u) represents a stencil that interpolates ∂zᶠᶜᶠ(i, j, k, grid, u)^2 from Face to Center in both x and z.

For instance, you could first create the model, define Ri and closure and then set! the model with this new closure and auxiliary field

Yes, your question makes a lot of sense! The reason we can't do that is because type of any of the properties of model cannot change. This means we need to know the type of model.closure before we build it. Another way to say this is that HydrostaticFreeSurfaceModel is concretely typed, which means that functions of model can be optimized based on the types of every property of HydrostaticFreeSurfaceModel. This gives us performance advantages / reduces overhead, but also creates some limitations like: we have to define a closure before building model. Here's a simple example:

julia> mutable struct Test{T}
       a :: T
       end

julia> t = Test(1.0)
Test{Float64}(1.0)

julia> t.a = 2.0
2.0

julia> t.a = "hi"
ERROR: MethodError: Cannot `convert` an object of type String to an object of type Float64
Closest candidates are:
  convert(::Type{T}, ::T) where T<:Number at number.jl:6
  convert(::Type{T}, ::Number) where T<:Number at number.jl:7
  convert(::Type{T}, ::Base.TwicePrecision) where T<:Number at twiceprecision.jl:250
  ...
Stacktrace:
 [1] setproperty!(x::Test{Float64}, f::Symbol, v::String)
   @ Base ./Base.jl:34
 [2] top-level scope
   @ REPL[6]:1

Here, t is Test{Float64}. We can change a as long as the new value can be converted to Float64 via convert(Float64, new_value). Otherwise we can't and we just have to write s = Test("hi") instead.

Another possible solution is to build a "dummy" closure with the correct type, and then replace it with the "real" closure after constructing model by writing model.closure = real_closure. It could be fun to try that and see if the code is simpler. I suspect it might not be but it's worth a shot (you might learn something).

Thinking on a more general way, it would be awesome to give a general function that could depend on any field in the model, including auxiliary. I am thinking that this way it would be easier for users to create and implement their own parameterizations.

I guess that's sort of what we're doing here. The main issue is reducing boilerplate / making this as easy as possible. It'd be nice to write parameterizations with more abstract syntax rather than having to write a low-level kernel function.

[glwagner]

Also you might want to make the diffusivity VerticalScalarDiffusivity rather than regular ScalarDiffusivity since I guess PP is just a vertical diffusion!

[glwagner]

Thinking on a more general way, it would be awesome to give a general function that could depend on any field in the model, including auxiliary.

Thinking about this more, I think we could "easily" (with some code changes") support a form similar to the "discrete form" forcing functions: i, j, k, grid, clock, model_fields and optional parameters. We'd probably need something like that to support parameter estimation / uncertainty quantification of user/script-defined closures with OceanLearning.jl so I think it's worth pondering on this for a bit.

It'd be more of a challenge to support a form like ContinuousForcing and allow derivatives of fields in user-defined functions (again possible, just a bit more complex to solve).

[glwagner]

This discussion is related to PR #2580