Transforms & Spectral

SHTnsKit configuration, MPI/PencilArrays parallel layout, and the spectral operator toolbox (energy spectra, rotations, filters, horizontal Laplacian).

Parallel And Transform API

GeoDynamo.SHTnsKitConfigType
SHTnsKitConfig <: AbstractSHTnsConfig

Main configuration structure for spherical harmonic transforms using SHTnsKit. This struct encapsulates all parameters needed for transforms and parallelization.

Fields

Core SHTnsKit Configuration

  • sht_config: The underlying SHTnsKit.SHTConfig object

Grid Parameters

  • nlat::Int: Number of latitude (theta) points (Gauss-Legendre grid)
  • nlon::Int: Number of longitude (phi) points (uniform grid)
  • lmax::Int: Maximum spherical harmonic degree
  • mmax::Int: Maximum spherical harmonic order (≤ lmax)
  • nlm::Int: Total number of (l,m) spectral mode pairs

Parallelization Infrastructure

  • pencils: Collection of PencilArrays Pencil objects for different data orientations (:theta, :phi, :r, :spec, :mixed)
  • fft_plans: Precomputed FFTW plans for FFT operations
  • transpose_plans: Plans for data redistribution between pencils

Auxiliary Data

  • memory_estimate::String: Human-readable memory usage estimate
  • l_values::Vector{Int}: Spherical harmonic degree for each spectral index
  • m_values::Vector{Int}: Spherical harmonic order for each spectral index
  • theta_grid::Vector{Float64}: Latitude values (Gauss-Legendre nodes)
  • phi_grid::Vector{Float64}: Longitude values (uniform spacing)
  • gauss_weights::Vector{Float64}: Gauss-Legendre quadrature weights

Internal

  • _buffers::SHTnsBuffers: Reusable work arrays to reduce allocations

Usage

config = create_shtnskit_config(lmax=32, mmax=32, nlat=64, nlon=128)
source
GeoDynamo.create_shtnskit_configFunction
create_shtnskit_config(; lmax, mmax, nlat, nlon, nr, optimize_decomp) -> SHTnsKitConfig

Create and initialize a complete SHTnsKit configuration for spherical harmonic transforms with MPI parallelization.

Keyword Arguments

  • lmax::Int: Maximum spherical harmonic degree (required)
  • mmax::Int=lmax: Maximum spherical harmonic order (≤ lmax)
  • nlat::Int: Number of latitude points. Must be ≥ lmax+1 for numerical accuracy. Defaults to max(lmax+2, 64)
  • nlon::Int: Number of longitude points. Must be ≥ 2*mmax+1 for alias-free transforms. Powers of 2 are preferred for FFT efficiency.
  • nr::Int: Number of radial points (required)
  • optimize_decomp::Bool=true: Whether to optimize MPI process topology

Returns

  • SHTnsKitConfig: Fully initialized configuration ready for transforms

Algorithm

  1. Create base SHTnsKit configuration with Gauss-Legendre grid
  2. Set up pencil decomposition for MPI parallelization
  3. Create FFT plans for longitude transforms
  4. Create transpose plans for data redistribution
  5. Initialize grid coordinates and quadrature weights

Example

# Create config for lmax=63 simulation
config = create_shtnskit_config(lmax=63, nlat=96, nlon=192)
source
GeoDynamo.shtnskit_spectral_to_physical!Function
shtnskit_spectral_to_physical!(spec, phys)

Transform spectral coefficients to physical space values (synthesis).

The Synthesis Operation

Given spherical harmonic coefficients fl^m, compute the physical field: f(θ,φ) = Σ{l,m} fl^m × Yl^m(θ,φ)

where Y_l^m are the spherical harmonics.

Implementation

Delegates to scalar_spectral_to_physical! (the Phase-2 r×θ θ-subcommunicator path), which works correctly under all process-grid configurations including r-distribution.

Arguments

  • spec::SHTnsSpecField: Source spectral field with coefficients
  • phys::SHTnsPhysField: Destination physical field (modified in-place)

Side Effects

Modifies phys.data with the synthesized field values

source
GeoDynamo.shtnskit_physical_to_spectral!Function
shtnskit_physical_to_spectral!(phys, spec)

Transform physical space values to spectral coefficients (analysis).

The Analysis Operation

Given physical field values f(θ,φ), compute spherical harmonic coefficients: fl^m = ∫∫ f(θ,φ) × Yl^m*(θ,φ) sin(θ) dθ dφ

The integral is computed numerically using:

  • Gauss-Legendre quadrature for the θ integral
  • FFT for the φ integral (exploiting periodicity)

Implementation

Delegates to scalar_physical_to_spectral! (the Phase-2 r×θ θ-subcommunicator path), which works correctly under all process-grid configurations including r-distribution.

Arguments

  • phys::SHTnsPhysField: Source physical field values
  • spec::SHTnsSpecField: Destination spectral field (modified in-place)

Side Effects

Modifies spec.data_real and spec.data_imag with the computed coefficients

source
GeoDynamo.shtnskit_vector_synthesis!Function
shtnskit_vector_synthesis!(tor_spec::SHTnsSpecField{T},
                          pol_spec::SHTnsSpecField{T},
                          vec_phys::SHTnsVectorField{T};
                          domain::Union{RadialDomain,Nothing}=nothing) where T

Vector synthesis using SHTnsKit spheroidal-toroidal decomposition with PencilArrays.

Toroidal-Poloidal Decomposition

For a solenoidal vector field v (∇·v = 0): v = ∇×(T r̂) + ∇×∇×(P r̂)

where T = toroidal scalar, P = poloidal scalar.

In spherical components: vr = l(l+1)/r * P * Ylm (from poloidal only) vθ, vφ from both T and P (computed by SHTnsKit.SHsphtortospat)

CRITICAL: SHTnsKit.SHsphtortospat returns ONLY tangential components. The radial component v_r MUST be computed separately from the poloidal scalar.

Parameters

  • domain: Optional RadialDomain needed for computing vr with correct radial scaling. If not provided, vr will be set to zero (suitable for tests).
source
GeoDynamo.shtnskit_vector_analysis!Function
shtnskit_vector_analysis!(vec_phys::SHTnsVectorField{T},
                         tor_spec::SHTnsSpecField{T},
                         pol_spec::SHTnsSpecField{T}) where T

Vector analysis using SHTnsKit with PencilArrays.

Toroidal-Poloidal Analysis

Decomposes a 3D velocity field into toroidal and poloidal scalars.

For a solenoidal vector field v (∇·v = 0): v = ∇×(T r̂) + ∇×∇×(P r̂)

Mathematical Note on Analysis

SHTnsKit.spattoSHsphtor takes (vθ, vφ) and returns (P, T).

This is mathematically valid for solenoidal fields because:

  1. The solenoidal constraint ∇·v = 0 couples vr to (vθ, v_φ)
  2. The decomposition into T (rotational) and P (potential) parts is unique
  3. The radial component v_r = l(l+1)/r * P is implicitly determined

However, this assumes EXACT solenoidality. In numerical simulations with finite precision, v_r may not exactly satisfy ∇·v = 0.

Alternative: Use Full 3-Component Analysis

For better numerical accuracy, one could use: Qcoeffs = analysis(vr * r / l(l+1)) # Recover P from vr S, T = spattoSHsphtor(vθ, v_φ) # Decompose tangential

Then check: Qcoeffs ≈ Scoeffs (should match for solenoidal field)

Current implementation uses 2-component analysis which is standard practice for solenoidal MHD simulations.

source
GeoDynamo.batch_shtnskit_transforms!Function
batch_shtnskit_transforms!(specs::Vector{SHTnsSpecField{T}},
                          physs::Vector{SHTnsPhysField{T}}) where T

Batch process multiple transforms using SHTnsKit with PencilArrays.

MPI Safety

This function processes transforms sequentially to avoid calling MPI collectives (Barrier) from multiple threads simultaneously. MPI collective operations must be called from the same thread on all processes to avoid deadlock.

Note: The individual transforms themselves are still efficient as SHTnsKit performs optimized Legendre transforms and FFTs internally.

source
GeoDynamo.set_shtnskit_threadsFunction
set_shtnskit_threads(num_threads::Int)

Configure the number of threads used by SHTnsKit transforms. Uses SHTnsKit.shtnsusethreads when available.

source
GeoDynamo.get_commFunction
get_comm()

Get MPI communicator, initializing MPI if needed. Thread-safe via double-checked locking with _MPI_INIT_LOCK.

source
GeoDynamo.create_pencil_topologyFunction
create_pencil_topology(shtns_config; nr, optimize=true)

Create enhanced pencil decomposition for SHTns grids. Phase 2 (r×θ 2D topology): reads the process grid from GEODYNAMOPROCGRID (θ_ranks×r_ranks), giving θ-distributed / φ-local / r-distributed physical pencils and a theta_phys 2D prototype built on the θ-subcommunicator. The optimize keyword is accepted for API compatibility. Accepts an object with fields nlat, nlon, nlm, lmax, and mmax (e.g., SHTnsKitConfig).

source
GeoDynamo.create_transpose_plansFunction
create_transpose_plans(pencils)

Create enhanced transpose plans between pencil orientations. Includes caching and communication optimization.

source
GeoDynamo.print_pencil_axesFunction
print_pencil_axes(pencils)

Print the axes_local tuple for each pencil, showing the local index ranges for all three axes. This helps verify which axes are distributed (those with nontrivial subranges across ranks) and which axis is contiguous locally.

source
GeoDynamo.validate_radial_distributionFunction
validate_radial_distribution(pencils; warn_uneven::Bool=true, strict::Bool=true) -> Bool

Validate that radial dimension has compatible distribution across all pencils.

MPI Synchronization Requirement

The SHTnsKit transforms use MPI.Allreduce inside per-radial-level loops. All processes must have the SAME number of local radial levels, otherwise processes will enter/exit the loop at different times causing MPI DEADLOCK.

Arguments

  • pencils: Named tuple of pencil configurations
  • warn_uneven: If true, emit warning for uneven distribution
  • strict: If true (default), throw an error instead of just warning

Returns

true if distribution is valid (all processes have same local radial count). false if there's a potential synchronization issue.

Default Behavior

Strict mode is enabled by default to prevent MPI deadlock in production runs. Set strict=false only for debugging purposes.

Example

# Recommended: use default strict mode
validate_radial_distribution(pencils)

# For debugging only: disable strict mode
validate_radial_distribution(pencils; strict=false)
source

Advanced Spectral Tools API

GeoDynamo.compute_scalar_energy_spectrumFunction
compute_scalar_energy_spectrum(config::SHTnsKitConfig, alm::Matrix{ComplexF64}; real_field::Bool=true)

Compute the energy spectrum per spherical harmonic degree l using SHTnsKit v1.1.15.

Returns

Vector of length lmax+1 with energy at each degree l.

source
GeoDynamo.compute_vector_energy_spectrumFunction
compute_vector_energy_spectrum(config::SHTnsKitConfig, Slm::Matrix{ComplexF64}, Tlm::Matrix{ComplexF64}; real_field::Bool=true)

Compute the kinetic energy spectrum per spherical harmonic degree l for a vector field decomposed into spheroidal (Slm) and toroidal (Tlm) components.

Returns

Vector of length lmax+1 with kinetic energy at each degree l.

source
GeoDynamo.compute_total_scalar_energyFunction
compute_total_scalar_energy(config::SHTnsKitConfig, alm::Matrix{ComplexF64}; real_field::Bool=true)

Compute total energy of a scalar field from its spectral coefficients.

source
GeoDynamo.compute_total_vector_energyFunction
compute_total_vector_energy(config::SHTnsKitConfig, Slm::Matrix{ComplexF64}, Tlm::Matrix{ComplexF64}; real_field::Bool=true)

Compute total kinetic energy of a vector field from spheroidal/toroidal coefficients.

source
GeoDynamo.compute_enstrophyFunction
compute_enstrophy(config::SHTnsKitConfig, Tlm::Matrix{ComplexF64}; real_field::Bool=true)

Compute enstrophy (mean square vorticity) from toroidal coefficients. Enstrophy is related to the rotational part of the kinetic energy.

source
GeoDynamo.spectral_gradient!Function
spectral_gradient!(config::SHTnsKitConfig, Slm::Matrix{ComplexF64},
                   grad_theta::Matrix{Float64}, grad_phi::Matrix{Float64})

Compute the horizontal gradient of a scalar field in spectral space. Uses SHTnsKit.SHtograd_spat for efficient computation.

Arguments

  • Slm: Spectral coefficients of the scalar field
  • grad_theta: Output θ-component of gradient (modified in-place)
  • grad_phi: Output φ-component of gradient (modified in-place)
source
GeoDynamo.extract_divergence_coefficientsFunction
extract_divergence_coefficients(config::SHTnsKitConfig, Slm::Matrix{ComplexF64})

Extract divergence spectral coefficients from spheroidal potential. The divergence field in spectral space is related to Slm by the horizontal Laplacian.

Returns

Matrix of divergence coefficients.

source
GeoDynamo.extract_vorticity_coefficientsFunction
extract_vorticity_coefficients(config::SHTnsKitConfig, Tlm::Matrix{ComplexF64})

Extract vorticity spectral coefficients from toroidal potential. The vorticity field in spectral space is related to Tlm by the horizontal Laplacian.

Returns

Matrix of vorticity coefficients.

source
GeoDynamo.shtnskit_qst_to_spatial!Function
shtnskit_qst_to_spatial!(config::SHTnsKitConfig, Qlm, Slm, Tlm, vr, vtheta, vphi)

Convert QST spectral coefficients to full 3D spatial vector field using SHTnsKit v1.1.15.

This is more efficient than separate scalar + vector synthesis as it handles all three components in a single call.

Arguments

  • Qlm: Q (radial) spectral coefficients
  • Slm: S (spheroidal/poloidal) spectral coefficients
  • Tlm: T (toroidal) spectral coefficients
  • vr, vtheta, vphi: Output spatial components (modified in-place)
source
GeoDynamo.shtnskit_spatial_to_qst!Function
shtnskit_spatial_to_qst!(config::SHTnsKitConfig, vr, vtheta, vphi, Qlm, Slm, Tlm)

Convert full 3D spatial vector field to QST spectral coefficients using SHTnsKit v1.1.15.

Arguments

  • vr, vtheta, vphi: Input spatial components
  • Qlm: Output Q (radial) spectral coefficients (modified in-place)
  • Slm: Output S (spheroidal/poloidal) spectral coefficients (modified in-place)
  • Tlm: Output T (toroidal) spectral coefficients (modified in-place)
source
GeoDynamo.shtnskit_synthesis_inplace!Function
shtnskit_synthesis_inplace!(config::SHTnsKitConfig, alm::Matrix{ComplexF64},
                             f_out::Matrix{Float64})

In-place spectral-to-physical synthesis using SHTnsKit v1.1.15. Writes result directly to f_out, avoiding allocation of temporary arrays.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients (lmax+1 × mmax+1)
  • f_out: Output physical field (nlat × nlon), modified in-place
source
GeoDynamo.shtnskit_analysis_inplace!Function
shtnskit_analysis_inplace!(config::SHTnsKitConfig, f::Matrix{Float64},
                            alm_out::Matrix{ComplexF64})

In-place physical-to-spectral analysis using SHTnsKit v1.1.15. Writes result directly to alm_out, avoiding allocation of temporary arrays.

Arguments

  • config: SHTnsKit configuration
  • f: Input physical field (nlat × nlon)
  • alm_out: Output spectral coefficients (lmax+1 × mmax+1), modified in-place
source
GeoDynamo.rotate_field_z!Function
rotate_field_z!(config::SHTnsKitConfig, alm::Matrix{ComplexF64}, alpha::Real;
                alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Rotate a scalar field around the z-axis by angle alpha in spectral space. This is a pure phase rotation: alm[l,m] -> alm[l,m] * exp(-imalpha)

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients (modified in-place if alm_out is nothing)
  • alpha: Rotation angle in radians
  • alm_out: Optional output array (if nothing, modifies alm in-place)

Returns

The rotated coefficients (alm_out if provided, otherwise alm)

source
GeoDynamo.rotate_field_y!Function
rotate_field_y!(config::SHTnsKitConfig, alm::Matrix{ComplexF64}, beta::Real;
                alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Rotate a scalar field around the y-axis by angle beta in spectral space. Uses Wigner d-matrices (small Wigner rotation matrices).

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • beta: Rotation angle in radians
  • alm_out: Optional output array

Returns

The rotated coefficients

source
GeoDynamo.rotate_field_90y!Function
rotate_field_90y!(config::SHTnsKitConfig, alm::Matrix{ComplexF64};
                  alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Rotate a scalar field by 90 degrees around the y-axis. This is a special case with optimized Wigner d-matrix values.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • alm_out: Optional output array

Returns

The rotated coefficients

source
GeoDynamo.rotate_field_90x!Function
rotate_field_90x!(config::SHTnsKitConfig, alm::Matrix{ComplexF64};
                  alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Rotate a scalar field by 90 degrees around the x-axis. Equivalent to: Z(-π/2) * Y(π/2) * Z(π/2)

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • alm_out: Optional output array

Returns

The rotated coefficients

source
GeoDynamo.rotate_field_euler!Function
rotate_field_euler!(config::SHTnsKitConfig, alm::Matrix{ComplexF64},
                    alpha::Real, beta::Real, gamma::Real;
                    alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Rotate a scalar field by Euler angles (ZYZ convention) in spectral space. The rotation is: R = Rz(gamma) * Ry(beta) * Rz(alpha)

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • alpha, beta, gamma: Euler angles in radians (ZYZ convention)
  • alm_out: Optional output array

Returns

The rotated coefficients

source
GeoDynamo.apply_horizontal_laplacian!Function
apply_horizontal_laplacian!(config::SHTnsKitConfig, alm::Matrix{ComplexF64};
                             alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Apply the horizontal (angular) Laplacian to spectral coefficients. ∇²h Yl^m = -l(l+1) Y_l^m

This scales each coefficient by -l(l+1), which is the eigenvalue of the horizontal Laplacian on the unit sphere.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • alm_out: Optional output array (if nothing, modifies alm in-place)

Returns

The Laplacian-transformed coefficients

source
GeoDynamo.apply_inverse_horizontal_laplacian!Function
apply_inverse_horizontal_laplacian!(config::SHTnsKitConfig, alm::Matrix{ComplexF64};
                                     alm_out::Union{Matrix{ComplexF64},Nothing}=nothing,
                                     regularize_l0::Bool=true)

Apply the inverse horizontal Laplacian to spectral coefficients. This scales each coefficient by -1/(l(l+1)), which is useful for solving Poisson equations.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • alm_out: Optional output array
  • regularize_l0: If true, set l=0 mode to zero (since 1/0 is undefined)

Returns

The inverse-Laplacian-transformed coefficients

source
GeoDynamo.compute_horizontal_gradient_magnitudeFunction
compute_horizontal_gradient_magnitude(config::SHTnsKitConfig, alm::Matrix{ComplexF64})

Compute the magnitude of the horizontal gradient |∇h f|² in spectral space. |∇h f|² = l(l+1) |f_lm|² (summed over all modes)

This is useful for computing gradient energy or penalty terms.

Arguments

  • config: SHTnsKit configuration
  • alm: Spectral coefficients

Returns

Scalar value of the integrated horizontal gradient magnitude squared

source
GeoDynamo.apply_spectral_filter!Function
apply_spectral_filter!(config::SHTnsKitConfig, alm::Matrix{ComplexF64},
                       filter_func::Function;
                       alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Apply a custom spectral filter to coefficients. The filter function takes (l, m) and returns a scaling factor.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • filter_func: Function (l, m) -> scale_factor
  • alm_out: Optional output array

Example

# Exponential filter for dealiasing
exp_filter(l, m) = exp(-(l/lmax)^16)
apply_spectral_filter!(config, alm, exp_filter)
source
GeoDynamo.apply_exponential_filter!Function
apply_exponential_filter!(config::SHTnsKitConfig, alm::Matrix{ComplexF64};
                           order::Int=16, cutoff::Float64=0.65,
                           alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Apply an exponential spectral filter for dealiasing. filter(l) = exp(-α * (l/lmax)^order) where α is chosen so filter(cutoff*lmax) = 0.5

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • order: Filter order (higher = sharper cutoff, default 16)
  • cutoff: Fraction of lmax where filter = 0.5 (default 0.65)
  • alm_out: Optional output array
source
GeoDynamo.truncate_spectral_modes!Function
truncate_spectral_modes!(config::SHTnsKitConfig, alm::Matrix{ComplexF64},
                         lmax_new::Int, mmax_new::Int=lmax_new;
                         alm_out::Union{Matrix{ComplexF64},Nothing}=nothing)

Truncate spectral coefficients to a lower resolution. Sets all modes with l > lmaxnew or m > mmaxnew to zero.

Arguments

  • config: SHTnsKit configuration
  • alm: Input spectral coefficients
  • lmax_new: New maximum degree
  • mmax_new: New maximum order (default = lmax_new)
  • alm_out: Optional output array
source
GeoDynamo.index_to_lm_fastFunction
index_to_lm_fast(idx, config) -> (l, m)

Fast O(1) conversion from linear spectral index to (l, m) using precomputed tables.

Uses the lvalues and mvalues arrays stored in SHTnsKitConfig during initialization. This is significantly faster than index_to_lm_shtnskit for repeated lookups.

Arguments

  • idx: Linear spectral index (1-based)
  • config: SHTnsKitConfig containing precomputed lvalues and mvalues

Returns

Tuple (l, m) for the spherical harmonic degree and order.

source
GeoDynamo.build_lm_lookup_tablesFunction
build_lm_lookup_tables(lmax, mmax) -> (l_values, m_values)

Build precomputed lookup tables for converting linear indices to (l, m).

Returns

  • l_values: Vector where l_values[idx] gives the degree l for linear index idx
  • m_values: Vector where m_values[idx] gives the order m for linear index idx
source
GeoDynamo.get_cached_buffer!Function
get_cached_buffer!(create_func, config, key::Symbol)
get_cached_buffer!(config, key::Symbol) do ... end

Thread-safe accessor for buffer cache. Returns existing buffer if present, otherwise creates a new one using create_func() and caches it.

Note: The function parameter comes FIRST to support Julia's do block syntax. When using do block, Julia desugars it to pass the closure as the first argument.

Arguments

  • create_func: Zero-argument function to create buffer if not cached
  • config: SHTnsKitConfig object containing _buffers::SHTnsBuffers
  • key::Symbol: Key to look up (mapped to a field of SHTnsBuffers)

Returns

The cached or newly created buffer.

Example

buffer = get_cached_buffer!(config, :my_buffer) do
    zeros(Float64, nlat, nlon)
end
source
GeoDynamo.clear_buffer_cache!Function
clear_buffer_cache!(config)

Thread-safe clearing of all lazily-allocated cached buffers. Useful when changing configurations or to free memory. Does NOT clear sht_plan, solver_transform_workspace, transform_device, spatial_scratch, or fft_scratch (these are set at construction/initialization and should persist).

source