Time Stepping

first_projection_step!(S, G, par, plans; a, dealias_mask=nothing, workspace=nothing, N2_profile=nothing)

Forward Euler initialization step for the leapfrog time stepper.

Purpose

The leapfrog scheme requires values at two time levels (n and n-1). This function takes the initial state and advances it by one Forward Euler step, providing the needed second time level.

Algorithm

1. Compute tendencies at time n:

   • Advection of q and B by geostrophic flow
   • Wave refraction by vorticity
   • Vertical diffusion

2. Apply physics switches:

   • linear: Zero nonlinear advection
   • inviscid: Zero dissipation
   • passive_scalar: Zero dispersion and refraction
   • fixed_flow: Mean flow doesn't evolve

3. Forward Euler update: For each spectral mode:

   q^(n+1) = [q^n - dt × tendency_q + dt × diffusion] × exp(-λ_q × dt)
   B^(n+1) = [B^n - dt × tendency_B] × exp(-λ_B × dt)

   where λ is the hyperdiffusion factor.

4. Wave feedback (optional):

   q* = q - qʷ

5. Diagnostic inversions:

   • q → ψ (elliptic inversion)
   • B → A, C (YBJ+ inversion)
   • ψ → u, v (velocity computation)

Arguments

• S::State: State to advance (modified in place)
• G::Grid: Grid struct
• par::QGParams: Model parameters
• plans: FFT plans
• a: Elliptic coefficient array a_ell(z) = f²/N²
• dealias_mask: Optional 2/3 dealiasing mask (nx × ny)
• workspace: Optional pre-allocated workspace for 2D decomposition
• N2_profile: Optional N²(z) profile for vertical velocity computation

Returns

Modified state S at time n+1.

Fortran Correspondence

This matches the projection step in main_waqg.f90.

Example

# Initialize and run projection step
state = init_state(grid)
init_random_psi!(state, grid)
a = a_ell_ut(params, grid)
L = dealias_mask(grid)
first_projection_step!(state, grid, params, plans; a=a, dealias_mask=L)

leapfrog_step!(Snp1, Sn, Snm1, G, par, plans; a, dealias_mask=nothing, workspace=nothing, N2_profile=nothing)

Advance the solution by one leapfrog time step with Robert-Asselin filtering.

Algorithm

1. Compute tendencies at time n:

F_q^n = J(ψ^n, q^n) - νz∂²q^(n-1)/∂z²
F_B^n = J(ψ^n, B^n) + dispersion + refraction

2. Leapfrog update with integrating factors: For each spectral mode (k):

q^(n+1) = q^(n-1) × e^(-2λdt) + 2dt × [-J(ψ,q)^n + diff^n] × e^(-λdt)
B^(n+1) = B^(n-1) × e^(-2λdt) + 2dt × [-J(ψ,B)^n + dispersion + refraction] × e^(-λdt)

Note: All tendencies are evaluated at time n and scaled by e^(-λdt) for second-order accuracy.

3. Robert-Asselin filter:

q̃^n = q^n + γ(q^(n-1) - 2q^n + q^(n+1))
B̃^n = B^n + γ(B^(n-1) - 2B^n + B^(n+1))

The filtered values are stored in Sn (which becomes Snm1 after rotation).

4. Wave feedback (if enabled):

q*^(n+1) = q^(n+1) - qʷ^(n+1)

5. Diagnostic inversions:

• q^(n+1) → ψ^(n+1)
• B^(n+1) → A^(n+1), C^(n+1)
• ψ^(n+1) → u^(n+1), v^(n+1)

Arguments

• Snp1::State: State at time n+1 (output)
• Sn::State: State at time n (input, filter applied to Snm1)
• Snm1::State: State at time n-1 (input, receives filtered values)
• G::Grid: Grid struct
• par::QGParams: Model parameters
• plans: FFT plans
• a: Elliptic coefficient array
• dealias_mask: Optional dealiasing mask
• workspace: Optional pre-allocated workspace for 2D decomposition
• N2_profile: Optional N²(z) profile for vertical velocity computation

Returns

Modified Snp1 with solution at time n+1.

Time Level Management

After this call:

• Snp1 contains fields at n+1
• Sn contains filtered fields at n (becomes new n-1 after rotation)
• Snm1 is unchanged (will be overwritten after rotation)

Typical loop structure:

for iter in 1:nsteps
    leapfrog_step!(Snp1, Sn, Snm1, G, par, plans; a=a)
    # Rotate: Snm1 ← Sn, Sn ← Snp1
    Snm1, Sn, Snp1 = Sn, Snp1, Snm1
end

Fortran Correspondence

This matches the main leapfrog loop in main_waqg.f90.

Example

# After projection step, run leapfrog
for iter in 1:1000
    leapfrog_step!(Snp1, Sn, Snm1, grid, params, plans; a=a, dealias_mask=L)
    Snm1, Sn, Snp1 = Sn, Snp1, Snm1
end

convol_waqg!(nqk, nBRk, nBIk, u, v, qk, BRk, BIk, G, plans; Lmask=nothing)

Compute advection terms in divergence form, matching Fortran convol_waqg.

Mathematical Form

Uses the divergence form of the Jacobian:

J(ψ, q) = ∂(uq)/∂x + ∂(vq)/∂y

where u, v are the geostrophic velocities (in real space).

Output

• nqk: Ĵ(ψ, q) - advection of QGPV
• nBRk: Ĵ(ψ, BR) - advection of wave real part
• nBIk: Ĵ(ψ, BI) - advection of wave imaginary part

Arguments

• nqk, nBRk, nBIk: Output arrays (spectral)
• u, v: Real-space velocity arrays (precomputed)
• qk, BRk, BIk: Input fields (spectral)
• G::Grid: Grid struct
• plans: FFT plans
• Lmask: Dealiasing mask (true = keep mode, false = zero)

Algorithm

For each field χ ∈ {q, BR, BI}:

1. Transform χ̂ → χ (inverse FFT)
2. Compute uχ and vχ (pointwise in real space)
3. Transform back: (ûχ), (v̂χ)
4. Compute divergence: ikₓ(ûχ) + ikᵧ(v̂χ)
5. Apply dealiasing mask

Fortran Correspondence

This matches convol_waqg in derivatives.f90.

Note

The velocities u, v should be precomputed and passed in real space.

convol_waqg_q!(nqk, u, v, qk, G, plans; Lmask=nothing)

Compute advection of q using divergence form without splitting wave fields.

convol_waqg_L⁺A!(nL⁺Ak, u, v, L⁺Ak, G, plans; Lmask=nothing)

Compute advection of complex L⁺A directly (YBJ+ path).

refraction_waqg!(rBRk, rBIk, BRk, BIk, psik, G, plans; Lmask=nothing)

Compute wave refraction term: B × ζ where ζ = ∇²ψ is relative vorticity.

Physical Interpretation

Near-inertial waves are refracted by vorticity gradients:

• Anticyclones (ζ < 0): Wave focusing, amplitude increase
• Cyclones (ζ > 0): Wave defocusing, amplitude decrease

This is the "wave capture" mechanism that traps NIWs in anticyclonic eddies.

Mathematical Form

refraction = B × ζ

where ζ = ∇²ψ = -kₕ²ψ̂ in spectral space.

Output

• rBRk: Real part of refraction term (spectral)
• rBIk: Imaginary part of refraction term (spectral)

Algorithm

1. Compute ζ̂ = -kₕ²ψ̂ (spectral)
2. Transform ζ̂, B̂R, B̂I to real space
3. Compute products: rBR = ζ × BR, rBI = ζ × BI
4. Transform back and apply dealiasing

Fortran Correspondence

This matches refraction_waqg in derivatives.f90.

Example

refraction_waqg!(rBR, rBI, BR, BI, psi, grid, plans; Lmask=L)
# rBR, rBI now contain the refraction tendencies

refraction_waqg_L⁺A!(rL⁺Ak, L⁺Ak, ψₖ, G, plans; Lmask=nothing)

Compute wave refraction term ζ*L⁺A directly for complex L⁺A (YBJ+ path).

compute_qw!(qwk, BRk, BIk, par, G, plans; Lmask=nothing)

Compute wave feedback on mean flow: qʷ from wave field B.

Physical Interpretation

The wave feedback qʷ represents how near-inertial waves modify the quasi-geostrophic flow. This is a key component of wave-mean flow interaction in the QG-YBJ+ model.

Mathematical Form (Xie & Vanneste 2015)

For dimensional equations where B has velocity units [m/s]:

qʷ = (i/2f)J(B*, B) + (1/4f)∇²|B|²

where:

• B* is the complex conjugate of B
• J(B, B) = BₓBᵧ - B*ᵧBₓ is the Jacobian
• |B|² = BR² + BI² is the wave energy density

No W2F scaling is applied since B already has its actual dimensional amplitude.

Decomposition

Let B = BR + i×BI. Then:

• J(B*, B) = 2i(BRₓBIᵧ - BRᵧBIₓ) [purely imaginary]
• ∇²|B|² = ∇²(BR² + BI²)

The final qʷ is real-valued after combining terms.

Arguments

• qwk: Output array for q̂ʷ (spectral)
• BRk, BIk: Wave field components (spectral)
• par: QGParams
• G::Grid: Grid struct
• plans: FFT plans
• Lmask: Dealiasing mask (true = keep mode, false = zero)

Example

compute_qw!(qw, BR, BI, params, grid, plans; Lmask=L)
# qw now contains wave feedback term

compute_qw_complex!(qʷₖ, Bk, par, G, plans; Lmask=nothing)

Compute wave feedback directly from complex B without spectral BR/BI splitting.

dissipation_q_nv!(dqk, qok, par, G; workspace=nothing)

Compute vertical diffusion of q with Neumann boundary conditions.

Mathematical Form

D = νz ∂²q/∂z²

with ∂q/∂z = 0 at z = -Lz and z = 0.

Discretization

Interior points (1 < k < nz): D[k] = νz (q[k+1] - 2q[k] + q[k-1]) / dz²

Boundary points (Neumann): D[1] = νz (q[2] - q[1]) / dz² D[nz] = νz (q[nz-1] - q[nz]) / dz²

Arguments

• dqk: Output array for diffusion term
• qok: Input q field at time n-1 (for leapfrog)
• par: QGParams (for nuz coefficient)
• G::Grid: Grid struct
• workspace: Optional pre-allocated workspace for 2D decomposition

Note

This operates on spectral q but the vertical derivative is in physical space, so the operation is the same for each (kx, ky) mode.

Fortran Correspondence

This matches dissipation_q_nv in derivatives.f90.

int_factor(kx, ky, par; waves=false)

Compute hyperdiffusion integrating factor for given wavenumber.

Mathematical Background

The hyperdiffusion operator is:

D = -ν₁(-∇²)^n₁ - ν₂(-∇²)^n₂

In spectral space, this becomes multiplication by:

λ = ν₁|k|^(2n₁) + ν₂|k|^(2n₂)

The integrating factor for one time step is: exp(-λ×dt)

For efficiency, we return just λ×dt (the exponent).

Arguments

• kx, ky: Horizontal wavenumber components
• par: QGParams (contains ν₁, ν₂, n₁, n₂)
• waves::Bool: If true, use wave hyperdiffusion (nuh1w, ilap1w, etc.)

Returns

λ×dt = dt × [ν₁(kx² + ky²)^n₁ + ν₂(kx² + ky²)^n₂] = dt × [ν₁ kₕ^(2n₁) + ν₂ kₕ^(2n₂)]

Note: Uses isotropic form (kx² + ky²)^n for proper damping of diagonal modes.

Usage in Time Stepping

# After computing tendency
factor = exp(-int_factor(kx, ky, par))
q_new = factor * q_tendency

Fortran Correspondence

This matches the integrating factor computation in the main loop of main_waqg.f90.

Example

# Get integrating factor for wavenumber (3, 4)
lambda_dt = int_factor(3.0, 4.0, params)
factor = exp(-lambda_dt)  # Multiply solution by this

init_imex_workspace(S, G; nthreads=Threads.maxthreadid())

Initialize workspace for IMEX time stepping with threading support.

NOTE: All work arrays are pre-allocated here to avoid heap corruption from repeated allocation/deallocation in tight loops during time stepping. Per-thread workspaces are created for the tridiagonal solves to enable parallel processing of horizontal modes.

The workspace includes storage for Adams-Bashforth 2 (previous tendency) and Strang splitting (intermediate B* state).

imex_cn_step!(Snp1, Sn, G, par, plans, imex_ws; a, dealias_mask=nothing, workspace=nothing, N2_profile=nothing)

Second-order IMEX-CNAB time step for YBJ+ equation with Strang splitting for refraction.

Uses Strang splitting with Adams-Bashforth 2 for advection (waves and mean flow) and Crank-Nicolson for linear PV diffusion/hyperdiffusion:

1. Stage 1 (First Half-Refraction): B* = B^n × exp(-i·(dt/2)·ζ/2)
2. Stage 2 (IMEX-CNAB for Advection + Dispersion):
   • Advection (AB2): (3/2)N^n - (1/2)N^{n-1} where N = -J(ψ,L⁺A)
   • Dispersion (CN): (1/2)[L(B*) + L(B^{n+1})]
3. Stage 3 (Second Half-Refraction): B^{n+1} = B** × exp(-i·(dt/2)·ζ/2) using ψ^{n+1} predictor

This achieves second-order temporal accuracy through:

• Strang splitting (second-order) instead of Lie splitting (first-order)
• Adams-Bashforth 2 (second-order) for advection (q and B)
• Crank-Nicolson for dispersion (second-order)

Algorithm

1. Compute q advection at time n: Q^n = -J(ψ^n, q^n)
2. Apply first half-refraction: B* = B^n × exp(-i·(dt/2)·ζ/2)
3. Compute B advection using B: N^n = -J(ψ^n, B)
4. For each spectral mode (kx, ky) solve IMEX-CNAB for B: a. Compute A* = (L⁺)⁻¹B* (essential for IMEX-CN consistency!) b. Build RHS = B* + (dt/2)·i·αdisp·kₕ²·A* + (3dt/2)·N^n - (dt/2)·N^{n-1} c. Solve modified elliptic: (L⁺ - β)·A^{n+1} = RHS where β = (dt/2)·i·αdisp·kₕ² d. Recover B** = RHS + β·A^{n+1}
5. Update q with CNAB: (I - dt/2·L)·q^{n+1} = (I + dt/2·L)·q^n + dt·[ (3/2)Q^n - (1/2)Q^{n-1} ] where L = νz∂zz - λ_h (vertical diffusion + hyperdiffusion)
6. Compute ψ^{n+1} predictor from q^{n+1} (and q^w predictor when enabled)
7. Apply second half-refraction using ψ^{n+1} predictor
8. Store tendencies for next step (AB2)

Arguments

• Snp1::State: State at time n+1 (output)
• Sn::State: State at time n (input)
• G::Grid: Grid
• par::QGParams: Parameters
• plans: FFT plans
• imex_ws::IMEXWorkspace: Pre-allocated workspace (stores N^{n-1} for AB2)
• a: Elliptic coefficient a_ell = f²/N²
• dealias_mask: Dealiasing mask
• workspace: Additional workspace for elliptic solvers
• N2_profile: N²(z) profile

Temporal Accuracy

• Overall: Second-order in time
• First step: Falls back to forward Euler for advection (AB2 bootstrap)

Stability

• Refraction: Exactly energy-preserving (integrating factor)
• Dispersion: Unconditionally stable (implicit CN)
• Advection CFL: dt < dx/U_max ≈ 1600s for this problem

first_imex_step!(S, G, par, plans, imex_ws; a, dealias_mask=nothing, workspace=nothing, N2_profile=nothing)

First-order forward Euler step to initialize IMEX time stepping.

This function:

1. Performs a first-order forward Euler step (same as firstprojectionstep!)
2. Initializes the AB2 state by computing and storing the q and B advection tendencies

After this function is called, subsequent calls to imexcnstep! will use second-order Adams-Bashforth 2 for advection.