Improves shear production rate calculations

[tomchor]

Closes https://github.com/tomchor/Oceanostics.jl/issues/103

CC @zwei961120

[tomchor]

@zwei961120 can you please check if this correctly addresses your issue?

About the pressure: I haven't ever done anything with the pressure to be honest. Do you mean that there should also be a separation between average and perturbation pressure?

[zwei961120]

Hi Tomas Chor,

I'm sorry for not getting back to you sooner. I have seen the improvements you have made to the shear production rate calculation, which is the same as my ideas. Thanks!

For the pressure flux term, in my opinion, when someone wants to do TKE budget analysis in the Z direction, the pressure field needs to be averaged horizontally (P) and calculate the separation (p - P), because in the horizontally averaged TKE budget equation, the pressure flux divergence term is (-∂z(w'p')), right? Of course, I don't know your thoughts on this point, feel free to discuss it with me.

Finally, do I need to leave comments on your improvements on GitHub, which might help more people?

Cheers, Zheng 发件人: Tomas Chor @.> 日期: 星期二, 2023年1月3日 23:27 收件人: tomchor/Oceanostics.jl @.> 抄送: Zheng WEI @.>, Mention @.> 主题: Re: [tomchor/Oceanostics.jl] Improves on shear production rate calculations (PR #104)

@zwei961120https://github.com/zwei961120 can you please check if this correctly addresses your issue?

About the pressure: I haven't ever done anything with the pressure to be honest. Do you mean that there should also be a separation between average and perturbation pressure?

— Reply to this email directly, view it on GitHubhttps://github.com/tomchor/Oceanostics.jl/pull/104#issuecomment-1369899677, or unsubscribehttps://github.com/notifications/unsubscribe-auth/A24HQMEZFI5MGPBSOATE3ZTWQRAMJANCNFSM6AAAAAATPXDP5Y. You are receiving this because you were mentioned.Message ID: @.***>

[tomchor]

For the pressure flux term, in my opinion, when someone wants to do TKE budget analysis in the Z direction, the pressure field needs to be averaged horizontally (P) and calculate the separation (p - P), because in the horizontally averaged TKE budget equation, the pressure flux divergence term is (-∂z(w'p')), right? Of course, I don't know your thoughts on this point, feel free to discuss it with me.

My thinking is the same. Can you please take a look at the latest modifications? I tried to change the notation to make it more intuitive. One way to use the new formulation is:

julia> u, v, w = model.velocities
NamedTuple with 3 Fields on 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo:
├── u: 4×4×4 Field{Face, Center, Center} on RectilinearGrid on CPU
├── v: 4×4×4 Field{Center, Face, Center} on RectilinearGrid on CPU
└── w: 4×4×5 Field{Center, Center, Face} on RectilinearGrid on CPU

julia> p = sum(model.pressures)
BinaryOperation at (Center, Center, Center)
├── grid: 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
└── tree: 
    + at (Center, Center, Center)
    ├── 4×4×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    └── 4×4×4 Field{Center, Center, Center} on RectilinearGrid on CPU

julia> W = Field(Average(w, dims=(2, 3)))
4×1×1 Field{Center, Nothing, Nothing} reduced over dims = (2, 3) on RectilinearGrid on CPU
├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (4, 1, 1)
├── grid: 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── operand: mean! over dims (2, 3) of 4×4×5 Field{Center, Center, Face} on RectilinearGrid on CPU
└── status: time=0.0

julia> P = Field(Average(p, dims=(2, 3)))
4×1×1 Field{Center, Nothing, Nothing} reduced over dims = (2, 3) on RectilinearGrid on CPU
├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (4, 1, 1)
├── grid: 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── operand: mean! over dims (2, 3) of BinaryOperation at (Center, Center, Center)
└── status: time=0.0

julia> wp = Oceanostics.TKEBudgetTerms.ZPressureRedistribution(model, w-W, p-P)
Derivative at (Center, Center, Center)
├── grid: 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
└── tree: 
    ∂zᶜᶜᶜ at (Center, Center, Center) via identity
    └── * at (Center, Center, Face)
        ├── - at (Center, Center, Face)
        │   ├── 4×4×5 Field{Center, Center, Face} on RectilinearGrid on CPU
        │   └── 4×1×1 Field{Center, Nothing, Nothing} reduced over dims = (2, 3) on RectilinearGrid on CPU
        └── - at (Center, Center, Center)
            ├── + at (Center, Center, Center)
            │   ├── 4×4×4 Field{Center, Center, Center} on RectilinearGrid on CPU
            │   └── 4×4×4 Field{Center, Center, Center} on RectilinearGrid on CPU
            └── 4×1×1 Field{Center, Nothing, Nothing} reduced over dims = (2, 3) on RectilinearGrid on CPU

Although I have a feeling that is complex enough that it won't compile on GPUs, in which case I'll re-write it using KernelFunctionOperations. Do you have any way to test that on GPUs? (I don't have access to one right now.)

Finally, do I need to leave comments on your improvements on GitHub, which might help more people?

I think this discussion being recorded here is already useful! That said, if you'd like to add something to the docstrings, or readme to make the usage of the package even clearer, it would be of course greatly appreciated!

The package unfortunately doesn't have docs yet, but if you're interested enough to help us create the docs (see issue https://github.com/tomchor/Oceanostics.jl/issues/99) that would of course also be greatly appreciated. Although creating the docs definitely isn't trivial work and would require some effort.

[tomchor]

@zwei961120 So the code in this repo isn't merged in the main branch of Oceanostics yet. So to test this code out, the best way is to install this specific branch of the package in your Project. This can be done as:

] add Oceanostics#tc/shearprod

Or using

using Pkg
Pkg.add("Oceanostics", rev="tc/shearprod")

[tomchor]

That said, your code here seems to be constructing things properly, although wp is never calculated, so we don't really know if it works on GPUs.

Here's a modified version of your code that (I think) would work:

using Oceananigans
using Oceananigans.Units: minute, minutes, hours
using Oceanostics
grid = RectilinearGrid(GPU();
                       size=(32, 32, 32),
                       extent=(128, 128, 64))

using Oceananigans.BuoyancyModels: g_Earth

amplitude = 0.8 # m
wavelength = 60  # m
wavenumber = 2π / wavelength # m⁻¹
frequency = sqrt(g_Earth * wavenumber) # s⁻¹

# The vertical scale over which the Stokes drift of a monochromatic surface wave
# decays away from the surface is `1/2wavenumber`, or
const vertical_scale = wavelength / 4π

# Stokes drift velocity at the surface
const Uˢ = amplitude^2 * wavenumber * frequency # m s⁻¹
uˢ(z) = Uˢ * exp(z / vertical_scale)
∂z_uˢ(z, t) = 1 / vertical_scale * Uˢ * exp(z / vertical_scale)
Qᵘ = -3.72e-5 # m² s⁻², surface kinematic momentum flux

u_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))
Qᵇ = 2.307e-8 # m² s⁻³, surface buoyancy flux
N² = 1.936e-5 # s⁻², initial and bottom buoyancy gradient

b_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ),
                                                bottom = GradientBoundaryCondition(N²))

Qᵇ = 2.307e-8 # m² s⁻³, surface buoyancy flux
N² = 1.936e-5 # s⁻², initial and bottom buoyancy gradient

b_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ),
                        bottom = GradientBoundaryCondition(N²))

coriolis = FPlane(f=1e-4) # s⁻¹
model = NonhydrostaticModel(; grid, coriolis,
                            advection = WENO(),
                            timestepper = :RungeKutta3,
                            tracers = :b,
                            buoyancy = BuoyancyTracer(),
                            closure = AnisotropicMinimumDissipation(),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            boundary_conditions = (u=u_boundary_conditions, b=b_boundary_conditions))
Ξ(z) = randn() * exp(z / 4)
initial_mixed_layer_depth = 33 # m
stratification(z) = z < - initial_mixed_layer_depth ? N² * z : N² * (-initial_mixed_layer_depth)

bᵢ(x, y, z) = stratification(z) + 1e-1 * Ξ(z) * N² * model.grid.Lz

u, v, w = model.velocities
p = sum(model.pressures)
W = Field(Average(w, dims=(2, 3)))
P = Field(Average(p, dims=(2, 3)))
w_prime = w-W
p_prime = p-P

wp_operation = ZPressureRedistribution(model, w_prime, p_prime)
wp = Field(wp_operation)

compute!(wp) # This line should throw an error on GPUs if my code is too complex

The code, in fact, can be made much shorter since we really only need w and p. It doesn't really matter what the ICs are, etc. So a "better" (i.e. shorter) test code would be

using Oceananigans
using Oceanostics
grid = RectilinearGrid(GPU();
                       size=(32, 32, 32),
                       extent=(128, 128, 64))
model = NonhydrostaticModel(; grid)

u, v, w = model.velocities
p = sum(model.pressures)
W = Field(Average(w, dims=(2, 3)))
P = Field(Average(p, dims=(2, 3)))
w_prime = w-W
p_prime = p-P

wp_operation = ZPressureRedistribution(model, w_prime, p_prime)
wp = Field(wp_operation)

compute!(wp) # This line should throw an error on GPUs if my code is too complex

Note that this would only work after you install this specific branch in your project! (Like I showed you in the previous comment.)

Note also that I haven't tested the codes above. So it's possible I made a typo somewhere... :grimacing:

[tomchor]

Finally, I didn't undetrstand your comment about the averaging dimensions here. Do you mind linking the specific part of the code you're referring to?

You make comments on specific lines of code by clicking the plus sign in here when you hover your mouse above it.

[zwei961120]

Thanks so much for your suggestions and tips! Actually, I'm new to Julia and Github and not familiar with them, but I will keep learning😊

For the comment about the averaging dimensions, I think I’m not confused now after thinking about it, because I think how to calculate average and separation depends on specific problems and cases, thank you for your detailed reply!

[tomchor]

Okay, apparently, this compiles on GPUs!:

julia> compute!(wp)
32×32×32 Field{Center, Center, Center} on RectilinearGrid on GPU
├── grid: 32×32×32 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on GPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
├── operand: Derivative at (Center, Center, Center)
├── status: time=0.0
└── data: 38×38×38 OffsetArray(::CUDA.CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, -2:35, -2:35, -2:35) with eltype Float64 with indices -2:35×-2:35×-2:35
    └── max=0.0, min=0.0, mean=0.0

@glwagner are you okay with me merging this as is? Then when we have the u - ZeroField() = u shortcut implemented in Oceananigans I can change the 0s to ZeroField()s.

[glwagner]

Can you summarize the changes?

[tomchor]

Can you summarize the changes?

Sure. The main changes are that the pressure redistribution and shear production rates calculations don't automatically use model.velocities for the velocities, they take user-defined velocities for the calculations. It's a necessary change because, like @zwei961120 pointed out, this means that our turbulent fluctuation velocities can now be defined accordingly to how we defined the averages. Put simply, the calculation before wasn't strictly correct.

In addition to that, I also changed the name from ZShearProduction() to ZShearProductionRate() (and same for x and y), since the term is actually the rate of the shear production, so this is a more correct name. I improved the relevant tests a bit.

[glwagner]

Ah sorry can you be a little more specific? What was the interface before, and what is it now? To calculate shear production the user needs to provide both the fluctuation and the mean correct?

[tomchor]

Ah sorry can you be a little more specific? What was the interface before, and what is it now? To calculate shear production the user needs to provide both the fluctuation and the mean correct?

For the shear prod rate the interface used to be just

XSP = XShearProduction(model, u, v, w, U, V, W)

where it was up to the user to define u, v, w, U, V and W manually. Now there are two interfaces :+1:

u, v, w = model.velocities
U = Field(Average(u, dims=(2, 3)))
V = Field(Average(v, dims=(2, 3)))
W = Field(Average(w, dims=(2, 3)))

XSP = XShearProductionRate(model, u-U, v-V, w-W, U, V, W) # manual one

XSP = XShearProductionRate(model; U=U, V=V, W=W) # this should produce the same result as before

As for the pressure redistribution term, the old function was this one:

function XPressureRedistribution(model)
    u, v, w = model.velocities
    p = sum(model.pressures)
    return ∂x(u*p)

So as you can it wasn't 100% correct unless of the averages was zero (i.e. it wasn't multiplying the turbulent fluctuation of u with the turbulent fluctuation of p; it was multiplying both full variables). The new version only has one interface that can be used as

u, v, w = model.velocities
p = sum(model.pressures)
W = Field(Average(w, dims=(2, 3)))
P = Field(Average(p, dims=(2, 3)))
w_prime = w-W
p_prime = p-P

wp = ZPressureRedistribution(model, w_prime, p_prime)

We could add another method for XPressureRegistribution where you only pass model and the averages.

Hope this makes sense

[glwagner]

Great very clear. I think that's a good hierarchy of interfaces; defaulting to the fluctuations model.velocities.u - U makes sense.

[glwagner]

One note is that this approximation to the shear production term is strictly only valid in the limit of vanishing time step or stationary flow. We should also check that the spatial stencil is formulated correctly to reflect the discrete conservation of kinetic energy.

[tomchor]

One note is that this approximation to the shear production term is strictly only valid in the limit of vanishing time step or stationary flow. We should also check that the spatial stencil is formulated correctly to reflect the discrete conservation of kinetic energy.

@glwagner Are you okay if I merge this as is for now? It's already an improvement, and we can modify the formulation to conserve KE in a future PR. (I'm planning on making that effort later for several other diagnostics.)

[glwagner]

Go for it!