Physics Functions

invert_q_to_psi!(S, G; a, par=nothing, workspace=nothing)

Invert spectral QGPV q(kx,ky,z) to obtain streamfunction ψ(kx,ky,z).

Mathematical Problem

For each horizontal wavenumber (kₓ, kᵧ), solve the vertical ODE:

∂/∂z(a(z) ∂ψ/∂z) - kₕ² ψ = q

which expands via the product rule to: a(z)∂²ψ/∂z² + a'(z)∂ψ/∂z - kₕ²ψ = q. The staggered discretization captures a'∂ψ/∂z automatically.

Neumann boundary conditions ψ_z = 0 at top and bottom.

Arguments

• S::State: State struct containing q (input) and psi (output)
• G::Grid: Grid struct with wavenumbers and vertical coordinates
• a::AbstractVector: Elliptic coefficient a_ell(z) = f²/N²(z), length nz
• par: Optional QGParams for density weighting (defaults to unity weights)
• workspace: Optional z-pencil workspace arrays for 2D decomposition

Implementation Details

For 2D decomposition:

1. Transpose q from xy-pencil to z-pencil (z becomes local)
2. Perform tridiagonal solve on z-pencil data
3. Transpose ψ from z-pencil back to xy-pencil

The discrete system is tridiagonal with structure (for interior row k):

• Upper diagonal: du[k] = ak+1
• Diagonal: d[k] = -(a[k+1] + a[k]) - kₕ² dz²
• Lower diagonal: dl[k] = ak

where a[k] = f₀²/N²[k] is evaluated at the interface below cell k.

Fortran Correspondence

This matches psi_solver in elliptic.f90.

Mean Mode (kₕ=0) Handling

For the horizontal mean mode (kₓ=kᵧ=0), the equation reduces to: ∂/∂z(a(z) ∂ψ/∂z) = q

With Neumann boundary conditions (∂ψ/∂z=0 at both boundaries), this operator is singular: the homogeneous equation has the constant function as its null space. Consequently:

1. A solution exists only if ∫q dz = 0 (compatibility condition)
2. The solution is determined only up to an arbitrary constant

This implementation sets ψ=0 for kₕ=0 because:

• For periodic domains, the mean streamfunction doesn't affect velocities (u = -∂ψ/∂y, v = ∂ψ/∂x, both zero for constant ψ)
• Standard spectral QG codes typically ignore the barotropic mean
• Initial conditions and forcing are assumed to have zero horizontal mean

If your application requires tracking vertically-varying barotropic modes, you would need to solve the singular ODE with an additional constraint (e.g., ∫ψ dz = 0) to uniquely determine the solution.

Example

a_vec = a_ell_ut(params, G)  # Compute a_ell = f²/N²
invert_q_to_psi!(state, grid; a=a_vec)

invert_L⁺A_to_A!(S, G, par, a; workspace=nothing)

YBJ+ wave amplitude recovery: solve for A given L⁺A (the prognostic variable).

Operator Definitions (from PDF)

L  (YBJ):   L  = ∂/∂z(f²/N² ∂/∂z)                        [eq. (4)]
L⁺ (YBJ+):  L⁺ = L - k_h²/4                               [spectral space]

Key relation: L = L⁺ + kh²/4, so LA = L⁺A + (kh²/4)A

Mathematical Problem

For each horizontal wavenumber (kₓ, kᵧ), solve the L⁺ elliptic equation:

L⁺A = ∂/∂z(a(z) ∂A/∂z) - (kₕ²/4) A = (L⁺A)_input

which expands via the product rule to:

a(z) ∂²A/∂z² + a'(z) ∂A/∂z - (kₕ²/4) A = (L⁺A)_input

with Neumann boundary conditions A_z = 0 at top and bottom.

The staggered discretization using Si = (f/(N(zi)Δz))² at different interfaces automatically captures the a'(z)∂A/∂z term from variable stratification.

Arguments

• S::State: State containing L⁺A (input), A and C (output)
• G::Grid: Grid struct
• par: QGParams (for f0, N2 parameters)
• a::AbstractVector: Elliptic coefficient a_ell(z) = f²/N²(z)
• workspace: Optional z-pencil workspace for 2D decomposition

Output Fields

• S.A: Recovered wave amplitude A
• S.C: Vertical derivative C = ∂A/∂z (for wave velocity computation)

Wave Velocity Amplitude

After this function, LA can be computed as: LA = L⁺A + (kh²/4)A This is used for GLM particle advection (see computewave_displacement!).

Mean Mode (kₕ=0) Handling

For the horizontal mean mode (kₓ=kᵧ=0), the equation reduces to: LA = ∂/∂z(a(z) ∂A/∂z) = L⁺A (since L⁺ = L when k_h=0)

With Neumann boundary conditions (∂A/∂z=0 at both boundaries), this operator is singular - the constant function is in its null space. To select a unique solution, we:

1. Fix a gauge (A[1]=0) to obtain a particular solution.
2. Remove the vertical mean of A (adds a constant null-space mode).

This yields a well-defined, mean-zero A for kₕ=0 while preserving the original equation.

Fortran Correspondence

This matches A_solver_ybj_plus in elliptic.f90 (PDF Eq. 33-35).

invert_helmholtz!(dstk, rhs, G, par; a, b=zeros, scale_kh2=1.0, bot_bc=nothing, top_bc=nothing, workspace=nothing)

General vertical Helmholtz inversion for each horizontal wavenumber.

Mathematical Problem

Solve the ODE:

a(z) ∂²φ/∂z² + b(z) ∂φ/∂z - scale_kh2 × kₕ² φ = rhs

with Neumann boundary conditions (∂φ/∂z specified at boundaries).

Discretization (matches Fortran helmholtzdouble)

Uses a centered stencil where same a[k], b[k] apply to all diagonals at point k:

• Bottom (k=1): d = -a[1] - 0.5b[1]Δz - αkₕ²Δz², du = a[1] + 0.5b[1]Δz
• Interior: d = -2a[k] - αkₕ²Δz², du = a[k] + 0.5b[k]Δz, dl = a[k] - 0.5b[k]Δz
• Top (k=nz): d = -a[nz] + 0.5b[nz]Δz - αkₕ²Δz², dl = a[nz] - 0.5b[nz]Δz

Boundary flux terms are added to RHS:

• Bottom: rhs[1] += (a[1] - 0.5b[1]Δz) × Δz × bot_bc
• Top: rhs[nz] -= (a[nz] + 0.5b[nz]Δz) × Δz × top_bc

Arguments

• dstk: Output array (nz, nx, ny) for solution φ
• rhs: Right-hand side array (nz, nx, ny)
• G::Grid: Grid struct
• par: QGParams (currently unused, kept for API consistency)
• a::AbstractVector: Second derivative coefficient a(z), length nz
• b::AbstractVector: First derivative coefficient b(z), length nz (default zeros)
• scale_kh2::Real: Multiplier α for kₕ² term (default 1.0)
• bot_bc, top_bc: Optional boundary flux arrays (nx, ny) for non-zero Neumann BCs
• workspace: Optional z-pencil workspace for 2D decomposition

Fortran Correspondence

This matches helmholtzdouble in elliptic.f90 exactly.

Note

For 2D decomposition, boundary conditions are not yet supported and will trigger a warning.

jacobian_spectral!(dstk, phik, chik, G, plans; Lmask=nothing)

Compute the Jacobian J(φ, χ) = ∂φ/∂x ∂χ/∂y - ∂φ/∂y ∂χ/∂x using pseudo-spectral method.

Mathematical Definition

The Jacobian (also called Poisson bracket) is:

J(φ, χ) = ∂φ/∂x ∂χ/∂y - ∂φ/∂y ∂χ/∂x

In vector form: J(φ, χ) = ẑ · (∇φ × ∇χ)

Physical Interpretation

• J(ψ, q): Advection of PV by geostrophic flow
• J(ψ, B): Advection of wave envelope by mean flow
• The Jacobian conserves both integrals ∫φ and ∫χ

Algorithm

1. Compute spectral derivatives: φ̂ₓ = ikₓφ̂, φ̂ᵧ = ikᵧφ̂
2. Transform derivatives to physical space
3. Compute product: J = φₓχᵧ - φᵧχₓ (pointwise)
4. Transform result back to spectral space

Arguments

• dstk: Output array for Ĵ(φ, χ) in spectral space
• phik: φ̂ in spectral space (must be real field, i.e., Hermitian symmetric)
• chik: χ̂ in spectral space (must be real field, i.e., Hermitian symmetric)
• G::Grid: Grid with wavenumber arrays
• plans: FFT plans from plan_transforms!
• Lmask: Optional 2/3 dealiasing mask (true = keep mode, false = zero)

Important

This function assumes φ and χ are real-valued fields in physical space. For real fields, IFFT of spectral derivatives (imkφ̂) yields real results (up to roundoff), so the physical derivatives are extracted via real().

Example

# Compute J(ψ, q) for real fields ψ and q
jacobian_spectral!(Jpsi_q, psi_k, q_k, grid, plans)

compute_velocities!(S, G; plans=nothing, params=nothing, compute_w=true, use_ybj_w=false, N2_profile=nothing, workspace=nothing, dealias_mask=nothing)

Compute geostrophic velocities from the spectral streamfunction ψ̂.

Physical Equations

Horizontal velocities from geostrophic balance:

u = -∂ψ/∂y  →  û(k) = -i kᵧ ψ̂(k)
v =  ∂ψ/∂x  →  v̂(k) =  i kₓ ψ̂(k)

Vertical velocity from QG omega equation:

∇²w + (f²/N²) ∂²w/∂z² = (2f/N²) J(ψ_z, ∇²ψ)

or YBJ formulation:

w = -(f²/N²) [(∂A/∂x)_z - i(∂A/∂y)_z] + c.c.

Algorithm

1. Compute û = -i kᵧ ψ̂ and v̂ = i kₓ ψ̂ in spectral space
2. Transform to physical space via inverse FFT
3. Optionally solve omega equation or use YBJ formula for w

Arguments

• S::State: State with ψ (input) and u, v, w (output)
• G::Grid: Grid with wavenumbers kx, ky
• plans: FFT plans (auto-generated if nothing)
• params: Model parameters (for f₀, N²)
• compute_w::Bool: If true, compute vertical velocity
• use_ybj_w::Bool: If true, use YBJ formula instead of omega equation
• N2_profile::Vector: Optional N²(z) profile for vertical velocity computation
• workspace: Optional pre-allocated workspace for 2D decomposition
• dealias_mask: Optional 2D dealiasing mask for omega equation RHS (quadratic term). Should be the same mask used for other nonlinear terms (typically 2/3 rule).

Returns

Modified State with updated u, v, w fields.

Note

This computes ONLY QG velocities. For Lagrangian advection including wave effects, use compute_total_velocities! instead.

Fortran Correspondence

Matches compute_velo in derivatives.f90.

compute_vertical_velocity!(S, G, plans, params; N2_profile=nothing, workspace=nothing, dealias_mask=nothing)

Solve the QG omega equation for ageostrophic vertical velocity.

Physical Background

In quasi-geostrophic dynamics, the leading-order horizontal flow is non-divergent (∇·u_g = 0). Vertical motion arises from ageostrophic corrections that maintain thermal wind balance as the flow evolves.

The omega equation relates w to the horizontal flow:

N² ∇²w + f² ∂²w/∂z² = 2f J(ψ_z, ∇²ψ)

or equivalently (dividing by N²):

∇²w + (f²/N²) ∂²w/∂z² = (2f/N²) J(ψ_z, ∇²ψ)

where:

• Left side: 3D Laplacian (horizontal + stratification-weighted vertical)
• Right side: Jacobian forcing from vertical shear and vorticity
• f²/N² << 1: stratification suppresses vertical motion relative to horizontal

Physical Interpretation

The RHS forcing J(ψ_z, ∇²ψ) represents:

• Thermal wind tilting: vertical shear ψ_z interacting with vorticity ∇²ψ
• Frontogenesis/frontolysis: differential advection of temperature gradients

Strong w occurs at:

• Fronts (sharp density gradients)
• Edges of eddies (strong vorticity gradients)

Numerical Method

1. Compute RHS in spectral space via omegaeqnrhs!
2. For each horizontal wavenumber (kₓ, kᵧ):
   • Set up tridiagonal system in z
   • Solve using LAPACK gtsv! (O(nz) per wavenumber)
3. Transform w to physical space

Boundary Conditions

w = 0 at z = -Lz and z = 0 (rigid lid and bottom).

Arguments

• S::State: State with ψ (input) and w (output)
• G::Grid: Grid structure
• plans: FFT plans
• params: Model parameters (f₀)
• N2_profile::Vector: Optional N²(z) profile (default: constant N² = 1)
• workspace: Optional pre-allocated workspace for 2D decomposition
• dealias_mask: 2D dealiasing mask for omega equation RHS. If nothing (default), a standard 2/3-rule mask is computed automatically to avoid aliasing in the quadratic Jacobian term.

Fortran Correspondence

Matches omega equation solver in the Fortran implementation.

compute_total_velocities!(S, G; plans=nothing, params=nothing, compute_w=true, use_ybj_w=false, N2_profile=nothing, workspace=nothing, dealias_mask=nothing, include_wave_velocity=true)

Compute the TOTAL velocity field for Lagrangian particle advection.

Physical Background

In QG-YBJ+ dynamics, a particle is advected by:

1. Geostrophic flow: uQG = -∂ψ/∂y, vQG = ∂ψ/∂x
2. Wave velocity: From YBJ+ equation (1.2): u + iv = e^{-ift} LA
3. Wave-induced Stokes drift: Second-order drift from near-inertial waves

The total velocity is:

u_total = u_QG + u_wave + u_S
v_total = v_QG + v_wave + v_S
w_total = w_QG + w_S (from omega equation or YBJ, plus vertical Stokes drift)

Wave Velocity (Asselin & Young 2019, eq. 1.2)

The backrotated wave velocity is LA, where L = ∂z(f²/N²)∂z:

u_wave = Re(LA)
v_wave = Im(LA)

For YBJ+: B = L⁺A where L⁺ = L + (1/4)Δ, so LA = B - (1/4)ΔA. In spectral space: LA = B + (k_h²/4)A

Wave-Induced Stokes Drift

Following Wagner & Young (2016) equations (3.16a)-(3.20), the Stokes drift uses the full Jacobian formulation:

Horizontal Stokes drift (eq. 3.16a, 3.18):

J₀ = (LA)* ∂_{s*}(LA) - (f₀²/N²)(∂_{s*} A_z*) ∂_z(LA)
u_S = Im(J₀)/f₀
v_S = -Re(J₀)/f₀

where ∂{s*} = (1/2)(∂x + i∂_y).

Vertical Stokes drift (eq. 3.19-3.20):

K₀ = M*_z · M_{ss*} - M*_{s*} · M_{sz}   where M = (f₀²/N²)A_z
w_S = -2·Im(K₀)/f₀

with Mz = az·A_z + a·A{zz}*, M{ss} = (a/4)·ΔH(Az), M{s*} = a·(A{zs})*, M{sz} = az·A{zs} + a·A{zzs}, and a = f₀²/N².

Usage

For Lagrangian particle advection, always use this function rather than compute_velocities! to include wave effects.

Arguments

• S::State: State with ψ, A, B (input) and u, v, w (output)
• G::Grid: Grid structure
• plans: FFT plans
• params: Model parameters
• compute_w::Bool: If true, compute vertical velocity
• use_ybj_w::Bool: If true, use YBJ formula for w
• N2_profile::Vector: Optional N²(z) profile for vertical velocity computation
• workspace: Optional pre-allocated workspace for 2D decomposition
• dealias_mask: Optional 2D dealiasing mask for omega equation RHS
• include_wave_velocity::Bool: If true (default), include wave velocity Re(LA), Im(LA)

Returns

Modified State with total velocity fields u, v, w.

flow_kinetic_energy(u, v) -> KE

Compute domain-integrated kinetic energy of the geostrophic flow (simple version).

Physical Background

The kinetic energy of the balanced flow:

KE = (1/2) ∫∫∫ (u² + v²) dx dy dz

This is a key diagnostic for:

• Model stability (unbounded growth indicates numerical issues)
• Energy conservation/dissipation rate
• Turbulent cascade analysis

Returns

Total kinetic energy (domain sum, not mean) in nondimensional units.

Note

• This is NOT normalized by volume. For energy density, divide by nx×ny×nz.
• In MPI mode, this returns LOCAL energy. Use mpireducesum for global total.
• For physically accurate energy with dealiasing and density weighting, use flow_kinetic_energy_spectral instead.

wave_energy(L⁺A, A) -> (E_L⁺A, E_A)

Compute domain-integrated wave energy from both L⁺A and A fields (simple version).

Physical Background

Two measures of wave energy in the model:

1. Envelope energy E_L⁺A = Σ |L⁺A|²

   • Based on the evolved YBJ+ wave envelope
   • Directly available from prognostic variable

2. Amplitude energy E_A = Σ |A|²

   • Based on the recovered wave amplitude
   • More physically meaningful for wave energy flux

Use Cases

• Monitor total wave energy conservation/dissipation
• Compare EL⁺A and EA to verify L⁺A→A recovery
• Track energy exchange with mean flow

Arguments

• L⁺A::Array{Complex,3}: YBJ+ wave envelope (spectral or physical)
• A::Array{Complex,3}: Wave amplitude (spectral or physical)

Returns

Tuple (EL⁺A, EA) of domain-summed squared magnitudes.

Note

• These are domain SUMS, not means. For energy density, divide by grid volume.
• In MPI mode, this returns LOCAL energy. Use mpireducesum for global total.
• For physically accurate wave energies with dealiasing and density weighting, use wave_energy_spectral instead.

flow_kinetic_energy_spectral(uk, vk, G, par; Lmask=nothing) -> KE

Compute kinetic energy in spectral space with dealiasing and density weighting.

Physical Background (matches Fortran diagzentrum/energylinear)

The kinetic energy is computed as:

KE(z) = Σₖ L(kₓ,kᵧ) × (|uₖ|² + |vₖ|²) - 0.5 × (|u₀₀|² + |v₀₀|²)

The dealiasing correction subtracts half the kh=0 mode because:

• With 2/3 dealiasing: Σₖ (1/2)|u|² = Σₖ L|u|² - 0.5|u(0,0)|²

The total KE integrates over z with density weighting:

KE_total = (1/nz) Σᵢ ρₛ(zᵢ) × KE(zᵢ)

Algorithm

1. Loop over all spectral modes (kₓ, kᵧ, z) with dealiasing mask L
2. Accumulate |u|² + |v|² at each level
3. Apply dealiasing correction: subtract half the kh=0 mode
4. Weight by density ρₛ(z) and integrate (divide by nz)

Arguments

• uk, vk: Spectral velocity fields (complex)
• G::Grid: Grid structure
• par: QGParams (for density profiles)
• Lmask: Optional dealiasing mask (default: all modes included)

Returns

Total kinetic energy, normalized by nz, with density weighting.

Fortran Correspondence

Matches the kinetic energy computation in diag_zentrum (diagnostics.f90:127-161) and energy_linear (diagnostics.f90:3024-3107).

Note

In MPI mode, returns LOCAL energy. Use mpireducesum for global total.

flow_potential_energy_spectral(bk, G, par; Lmask=nothing) -> PE

Compute potential energy in spectral space with dealiasing and density weighting.

Physical Background

The potential energy from buoyancy variance:

PE(z) = Σₖ L(kₓ,kᵧ) × (a_ell × ρ₁/ρ₂) × |bₖ|² - 0.5 × correction

where a_ell = f²/N² is the elliptic coefficient.

For QG: b = ψ_z, so PE represents available potential energy from isopycnal tilting.

Arguments

• bk: Spectral buoyancy field (complex)
• G::Grid: Grid structure
• par: QGParams (for f0, N2 and density profiles)
• Lmask: Optional dealiasing mask

Returns

Total potential energy, normalized by nz, with density weighting.

Fortran Correspondence

Matches the potential energy computation in diag_zentrum (ps term).

wave_energy_spectral(BR, BI, AR, AI, CR, CI, G, par; Lmask=nothing) -> (WKE, WPE, WCE)

Compute physically accurate wave energies in spectral space with dealiasing.

Operator Definitions (from PDF)

L  (YBJ operator):   L  = ∂/∂z(f²/N² ∂/∂z)              [eq. (4)]
L⁺ (YBJ+ operator):  L⁺ = L - k_h²/4                     [spectral space]

Key relation: L = L⁺ + kh²/4, so LA = B + (kh²/4)A

Physical Background (matches YBJ+ paper)

Three components of wave energy:

1. Wave Kinetic Energy (WKE) - uses YBJ L operator (NOT L⁺): WKE = (1/2) × Σₖ |LAₖ|²

   where LA is computed using the YBJ L operator [eq. (4)]: L = ∂/∂z(f²/N² × ∂/∂z)

   So LA = ∂_z(a(z) × C) where a(z) = f²/N² and C = ∂A/∂z. This is discretized as: LA[k] = (a[k]×C[k] - a[k-1]×C[k-1]) / Δz

   Note: This LA is the wave velocity amplitude, same as used in GLM particle advection for computing wave displacement ξ.

2. Wave Potential Energy (WPE): WPE = Σₖ (0.5/(ρ₂×a_ell)) × kh² × (|CRₖ|² + |CIₖ|²)

   where C = ∂A/∂z and a_ell = f²/N². This represents the potential energy from vertical wave structure.

3. Wave Correction Energy (WCE): WCE = Σₖ (1/8) × (1/a_ell²) × kh⁴ × (|ARₖ|² + |AIₖ|²)

   Higher-order correction term from the YBJ+ formulation (related to L vs L⁺ difference).

Algorithm

1. Loop over all spectral modes with dealiasing mask L
2. Compute LA = ∂_z(a×C) using finite differences for each (i,j) mode
3. Accumulate |LA|², kh²|C|²/(ρ₂×aell), kh⁴|A|²/(8×aell²)
4. Apply dealiasing correction: subtract half the kh=0 mode from WKE
5. Integrate over z (sum local, divide by nz)

Arguments

• BR, BI: Real and imaginary parts of wave envelope B (spectral)
• AR, AI: Real and imaginary parts of wave amplitude A (spectral)
• CR, CI: Real and imaginary parts of C = ∂A/∂z (spectral)
• G::Grid: Grid structure
• par: QGParams (for f0, N2)
• Lmask: Optional dealiasing mask

Returns

Tuple (WKE, WPE, WCE) of wave energy components, normalized by nz.

Fortran Correspondence

Matches wave_energy subroutine in diagnostics.f90 (lines 647-743).

Note

In MPI mode, returns LOCAL energy. Use mpireducesum for global totals.

flow_kinetic_energy_global(u, v, mpi_config=nothing) -> KE

Compute GLOBAL domain-integrated kinetic energy across all MPI processes.

Arguments

• u, v: Velocity arrays (local portion in MPI mode)
• mpi_config: MPI configuration (nothing for serial mode)

Returns

Global kinetic energy (sum across all processes).

Example

# Serial mode
KE = flow_kinetic_energy_global(state.u, state.v)

# MPI mode
KE = flow_kinetic_energy_global(state.u, state.v, mpi_config)

wave_energy_global(L⁺A, A, mpi_config=nothing) -> (E_L⁺A, E_A)

Compute GLOBAL wave energy across all MPI processes.

Arguments

• L⁺A, A: YBJ+ wave envelope and amplitude arrays (local portion in MPI mode)
• mpi_config: MPI configuration (nothing for serial mode)

Returns

Tuple (EL⁺A, EA) of global summed squared magnitudes.

Example

# Serial mode
EL⁺A, EA = wave_energy_global(state.L⁺A, state.A)

# MPI mode
EL⁺A, EA = wave_energy_global(state.L⁺A, state.A, mpi_config)

flow_kinetic_energy_spectral_global(uk, vk, G, par; Lmask=nothing, mpi_config=nothing) -> KE

Compute GLOBAL kinetic energy in spectral space across all MPI processes.

Arguments

• uk, vk: Spectral velocity fields (local portion in MPI mode)
• G::Grid: Grid structure
• par: QGParams
• Lmask: Optional dealiasing mask
• mpi_config: MPI configuration (nothing for serial mode)

Returns

Global kinetic energy with dealiasing and density weighting.

flow_potential_energy_spectral_global(bk, G, par; Lmask=nothing, mpi_config=nothing) -> PE

Compute GLOBAL potential energy in spectral space across all MPI processes.

Arguments

• bk: Spectral buoyancy field (local portion in MPI mode)
• G::Grid: Grid structure
• par: QGParams
• Lmask: Optional dealiasing mask
• mpi_config: MPI configuration (nothing for serial mode)

Returns

Global potential energy with dealiasing and density weighting.

wave_energy_spectral_global(BR, BI, AR, AI, CR, CI, G, par; Lmask=nothing, mpi_config=nothing) -> (WKE, WPE, WCE)

Compute GLOBAL wave energies in spectral space across all MPI processes.

Arguments

• BR, BI, AR, AI, CR, CI: Spectral wave fields (local portions in MPI mode)
• G::Grid: Grid structure
• par: QGParams
• Lmask: Optional dealiasing mask
• mpi_config: MPI configuration (nothing for serial mode)

Returns

Tuple (WKE, WPE, WCE) of global wave energy components.

EnergyDiagnosticsManager

Manages separate energy diagnostic output files.

record_energies!(manager, time, wave_KE, wave_PE, wave_CE, mean_flow_KE, mean_flow_PE)

Record energy values to the internal time series.

write_all_energy_files!(manager::EnergyDiagnosticsManager)

Write all energy time series to their respective files.

omega_eqn_rhs!(rhs, psi, G, plans; Lmask=nothing, workspace=nothing)

Compute the RHS forcing for the QG omega equation.

Physical Background

The QG omega equation relates vertical velocity w to the horizontal flow:

N² ∇²w + f² ∂²w/∂z² = 2f J(ψ_z, ∇²ψ)

or equivalently (dividing by N²):

∇²w + (f²/N²) ∂²w/∂z² = (2f/N²) J(ψ_z, ∇²ψ)

This function computes 2 J(ψ_z, ∇²ψ). The solver then applies the (f/N²) scaling.

Physical Interpretation

The Jacobian J(ψ_z, ∇²ψ) represents:

• ψ_z: Vertical shear of streamfunction (related to thermal wind/buoyancy)
• ∇²ψ: Relative vorticity ζ
• J: Cross-gradient interaction

Strong RHS forcing occurs where:

• Fronts (large ψ_z) interact with vorticity gradients
• Eddies tilt isopycnals through differential advection

Numerical Method

1. Vertical derivative: ψ_z via forward finite difference

   ψ_z[k] = (ψ[k+1] - ψ[k]) / dz,  ψ_z[nz] = 0 (Neumann)

2. Spectral derivatives:

   • ∂ψz/∂x = i kₓ ψz
   • ∂ψz/∂y = i kᵧ ψz
   • ∂(∇²ψ)/∂x = -i kₓ kh² ψ_avg
   • ∂(∇²ψ)/∂y = -i kᵧ kh² ψ_avg

   where ψ_avg = (ψ[k+1] + ψ[k])/2 for staggered-grid consistency

3. Jacobian in physical space:

   J(ψ_z, ∇²ψ) = (∂ψ_z/∂x)(∂∇²ψ/∂y) - (∂ψ_z/∂y)(∂∇²ψ/∂x)

4. Transform back: FFT to get spectral RHS

Arguments

• rhs::Array{Complex,3}: Output RHS array (modified in-place)
• psi::Array{Complex,3}: Spectral streamfunction
• G::Grid: Grid structure
• plans: FFT plans
• Lmask: Optional dealiasing mask
• workspace: Optional pre-allocated workspace for 2D decomposition

Returns

Modified rhs array with the omega equation forcing.

Fortran Correspondence

Matches omega_eqn_rhs computation in the Fortran implementation.

fft_forward!(dst, src, P::Plans)

Compute horizontal forward FFT (complex-to-complex) for each z-plane.

Algorithm

Serial FFTW backend: Loops over z-slices and applies 2D FFT to each (x,y) plane.

Arguments

• dst: Destination array (spectral space)
• src: Source array (physical space)
• P::Plans: FFT plans

Returns

Modified dst array.

Note

For parallel execution with PencilArrays, the MPI support provides a separate fft_forward!(dst::PencilArray, src::PencilArray, plans::MPIPlans) method that handles distributed transforms automatically.

fft_backward!(dst, src, P::Plans)

Compute horizontal inverse FFT (complex-to-complex) for each z-plane.

Algorithm

Serial FFTW backend: Loops over z-slices and applies 2D inverse FFT to each (x,y) plane. FFTW.ifft is NORMALIZED (divides by N automatically).

Arguments

• dst: Destination array (physical space, normalized)
• src: Source array (spectral space)
• P::Plans: FFT plans

Returns

Modified dst array.

Note

For parallel execution with PencilArrays, the MPI support provides a separate fft_backward!(dst::PencilArray, src::PencilArray, plans::MPIPlans) method that uses ldiv! for normalized inverse transforms.

dealias_mask(G) -> Matrix{Bool}

Compute the 2/3-rule dealiasing mask for spectral space.

Physical Background

In pseudo-spectral methods, nonlinear terms are computed by:

1. Transform fields to physical space (inverse FFT)
2. Compute products in physical space
3. Transform back to spectral space (forward FFT)

The problem: A product of two fields with max wavenumber kmax produces wavenumbers up to 2×kmax. With finite resolution, these high-k components "fold back" (alias) onto resolved wavenumbers, causing errors.

The 2/3 Rule

To prevent aliasing from quadratic nonlinearities (e.g., u·∇q):

• Keep only wavenumbers |k| ≤ (2/3) × k_Nyquist
• Truncated modes: set to zero before computing nonlinear products
• Result: product wavenumbers stay within (2/3)×2 = (4/3) < k_Nyquist

This rule is exact for quadratic nonlinearities in 1D. For 2D with radial cutoff, it provides effective dealiasing.

Algorithm

Uses radial (isotropic) cutoff:

• k_max = min(nx, ny) / 3
• Keep mode (i,j) if sqrt(kx² + ky²) ≤ k_max
• More isotropic than rectangular truncation

Arguments

• G::Grid: Grid with dimensions nx, ny

Returns

Matrix{Bool} of size (nx, ny):

• true = keep this wavenumber
• false = truncate (set to zero)

Usage

mask = dealias_mask(G)
q_hat .*= mask  # Zero out aliased modes

Fortran Correspondence

Matches LL(i,j) mask in the Fortran implementation.

compute_hyperdiff_coeff(; dx, dy, dt, order=4, efold_steps=10, kmax_fraction=1.0)

Compute hyperdiffusion coefficient for given grid spacing and desired damping rate.

Mathematical Background

The n-th order hyperdiffusion operator is: -νn ∇^n q = -νn (kₓ² + kᵧ²)^(n/2) q̂ in spectral space

The damping rate at wavenumber k is: λ = ν_n × k^n

For the grid scale (kmax = π/Δx), we want the amplitude to decay by factor e^(-1) after `efoldsteps` time steps, giving: νn = 1 / (efoldsteps × dt × k_max^n)

Arguments

• dx, dy: Grid spacing in x and y [m]
• dt: Time step [s]
• order: Order of hyperdiffusion (4 = biharmonic, 6 = hyper-6, etc.)
• efold_steps: Number of time steps for e-folding at grid scale (default: 10)
• kmax_fraction: Fraction of Nyquist wavenumber to target (default: 1.0)

Returns

Hyperdiffusion coefficient ν_n with units [m^n / s]

Example

# For 4th order (biharmonic) hyperdiffusion on a 1km grid with dt=10s
ν₄ = compute_hyperdiff_coeff(dx=1e3, dy=1e3, dt=10.0, order=4, efold_steps=5)

# Use in default_params
par = default_params(nx=128, ny=128, nz=64, Lx=128e3, Ly=128e3, Lz=3000.0,
                     dt=10.0, nt=1000,
                     νₕ₁=ν₄, νₕ₂=0.0, ilap1=2, ilap2=2)  # Pure 4th order

Notes

• Smaller efold_steps → stronger damping (more dissipative)
• Larger efold_steps → weaker damping (less dissipative)
• For stability, typically use efold_steps ∈ [5, 20]
• The 4th order (biharmonic) is most common: order=4, ilap1=2

compute_hyperdiff_params(; nx, ny, Lx, Ly, dt, order=4, efold_steps=10)

Convenience function to compute hyperdiffusion coefficient from grid parameters.

Arguments

• nx, ny: Grid points in x and y
• Lx, Ly: Domain size in x and y [m]
• dt: Time step [s]
• order: Order of hyperdiffusion (4 = biharmonic, default)
• efold_steps: E-folding time in steps (default: 10)

Returns

Named tuple (ν₄=..., ilap=...) for use with default_params

Example

# Compute 4th order hyperdiffusion for Asselin example
hd = compute_hyperdiff_params(nx=128, ny=128, Lx=70e3, Ly=70e3, dt=10.0, efold_steps=5)

par = default_params(nx=128, ny=128, nz=64, Lx=70e3, Ly=70e3, Lz=3000.0,
                     dt=10.0, nt=1000,
                     νₕ₁=hd.ν, νₕ₂=0.0, ilap1=hd.ilap, ilap2=2)

sumL⁺A!(B, G; Lmask=nothing, workspace=nothing)

Remove the vertical mean from the wave envelope B at each horizontal wavenumber.

Physical Background

In the normal YBJ formulation, the wave envelope B is related to amplitude A by:

B = (∂²A/∂z²) / N²

Since ∂²A/∂z² must integrate to zero (boundary conditions), B should have zero vertical mean. This function enforces that constraint.

Algorithm

For each horizontal wavenumber (kₓ, kᵧ) within the dealiasing mask:

1. Compute vertical mean: B̄(kₓ,kᵧ) = (1/nz) Σₖ B(kₓ,kᵧ,k)
2. Subtract mean: B(kₓ,kᵧ,k) ← B(kₓ,kᵧ,k) - B̄

For wavenumbers outside the mask or kh² = 0, set B = 0.

Arguments

• B::Array{Complex,3}: Wave envelope (modified in-place)
• G::Grid: Grid structure with wavenumbers
• Lmask: Optional dealiasing mask (default: all modes kept)
• workspace: Optional pre-allocated workspace for 2D decomposition

Returns

Modified B array with zero vertical mean at each (kₓ, kᵧ).

Fortran Correspondence

Matches sumB in derivatives.f90.

compute_sigma(par, G, nBRk, nBIk, rBRk, rBIk; Lmask=nothing, workspace=nothing, N2_profile=nothing) -> sigma

Compute the sigma constraint for normal YBJ A recovery.

Physical Background

When recovering A from B via vertical integration, we need to determine the vertical mean of A. The sigma parameter provides this constraint from the nonlinear forcing terms.

Mathematical Formula

For each horizontal wavenumber (kₓ, kᵧ):

σ(kₓ,kᵧ) = Σₖ [(rBRk + 2·nBIk)/kh² + i(rBIk - 2·nBRk)/kh²]

where:

• nBRk, nBIk: Real and imaginary parts of nonlinear advection term
• rBRk, rBIk: Real and imaginary parts of refraction term
• kh² = kₓ² + kᵧ²

Arguments

• par::QGParams: Model parameters
• G::Grid: Grid with wavenumbers
• nBRk, nBIk: Real/imaginary parts of advection forcing
• rBRk, rBIk: Real/imaginary parts of refraction forcing
• Lmask: Optional dealiasing mask
• workspace: Optional pre-allocated workspace for 2D decomposition
• N2_profile: Optional N²(z) profile for consistent stratification physics. If not provided, uses N2_ut(par, G) based on par.stratification.

Returns

2D complex array sigma(nxlocal, nylocal) with the constraint values.

Fortran Correspondence

Matches compute_sigma in derivatives.f90.

Note

In MPI mode with 2D decomposition, this requires z to be fully local. Transpose operations are handled internally if needed.

compute_A!(A, C, BRk, BIk, sigma, par, G; Lmask=nothing, workspace=nothing, N2_profile=nothing)

Recover wave amplitude A from envelope B using normal YBJ vertical integration.

Physical Background

In normal YBJ, B and A are related by:

B = (∂²A/∂z²) / N²

To recover A from B, we integrate twice using the N² weight:

1. First integral: ∂A/∂z = ∫ B · N² dz + c₁
2. Second integral: A = ∫∫ B · N² dz² + c₁z + c₂

The constants are determined by:

• Boundary condition: ∂A/∂z = 0 at top (Neumann)
• Mean constraint: ∫A dz = σ (from sigma)

Algorithm

For each horizontal wavenumber (kₓ, kᵧ):

Stage 1: Cumulative Integration

Ã[1] = 0
Ã[k] = Ã[k-1] + (Σⱼ₌₁ᵏ⁻¹ B[j]) × N²[k-1] × dz²

Stage 2: Apply Sigma Constraint

sumA = Σₖ Ã[k]
adj = (σ - sumA) / nz
A[k] = Ã[k] + adj   # Enforce ∫A = σ

Stage 3: Compute Vertical Derivative

C[k] = (A[k+1] - A[k]) / dz   # Forward difference
C[nz] = 0                      # Neumann BC at top

Arguments

• A::Array{Complex,3}: Output wave amplitude (modified in-place)
• C::Array{Complex,3}: Output vertical derivative A_z (modified in-place)
• BRk, BIk: Real/imaginary parts of wave envelope B
• sigma::Array{Complex,2}: Sigma constraint from compute_sigma
• par::QGParams: Model parameters
• G::Grid: Grid structure
• Lmask: Optional dealiasing mask
• workspace: Optional pre-allocated workspace for 2D decomposition
• N2_profile: Optional N²(z) profile for variable stratification. If not provided, uses N2_ut(par, G).

Returns

Tuple (A, C) with recovered amplitude and its vertical derivative.

Fortran Correspondence

Matches compute_A in derivatives.f90.

Note

This is the NORMAL YBJ recovery method. For YBJ+, use invert_L⁺A_to_A! instead, which solves the full L⁺A = B elliptic problem via tridiagonal solve.