Internals

_BUFFER_CACHE_LOCK

Global ReentrantLock for thread-safe access to SHTnsKitConfig buffer caches. All access to config._buffers should be protected by this lock.

Abstract type for scalar fields (temperature, composition, etc.)

AnalyticIC(pattern::Symbol; amplitude=1.0, parameters...)

Set a field to an analytical pattern defined in InitialConditions.

Supported patterns (see set_analytical_initial_conditions! for details):

• :conductive – conductive temperature profile
• :dipole – dipolar magnetic field
• :convective – small convective velocity perturbations
• :stratified – stratified composition profile
• :hot_blob – hot thermal blob
• :uniform_field – uniform magnetic field
• :blob – compositional blob

Any extra keyword arguments are forwarded verbatim to the underlying set_analytical_initial_conditions! call.

CB3()

Compatibility alias for RungeKutta3.

CNAB2(; theta=0.5)

Crank-Nicolson Adams-Bashforth second-order timestepper.

Callback(f; schedule)
Callback(f, schedule)

Wraps an arbitrary callable f with a firing schedule. When the callback fires, f(sim) is called with the current Simulation object.

CheckpointWriter(path; schedule)

Schedule-driven writer that writes a restart/checkpoint file to path whenever schedule fires.

EAB2(; krylov_dimension=20, tolerance=1e-8)

Compatibility alias for ExponentialAdamsBashforth2.

ERK2()

Compatibility alias for ExponentialRungeKutta2.

GeoDynamo.ERK2Cache(args...)

Legacy constructor that builds a Float64 ERK2 cache when no element type is specified explicitly.

GeoDynamo.ERK2Cache{T}(...)

Compatibility constructor for the public ERK2 cache type.

The solver now stores stage caches as ERK2StageCache; this constructor keeps older call sites that instantiate GeoDynamo.ERK2Cache directly working.

ETD(; krylov_dimension=20, tolerance=1e-8)

Exponential time differencing timestepper.

EnergyDiagnostics(; schedule)

Logs the kinetic energy (and magnetic energy, when a magnetic field is present) at each scheduled interval, computed from the model's spectral coefficients.

EnergyTracker

In-memory accumulator for optional diagnostic energy histories.

The NetCDF writer currently uses the global ENERGY_TRACKER for lightweight energy bookkeeping in tests and diagnostics helpers.

ExponentialAdamsBashforth2(; krylov_dimension=20, tolerance=1e-8)

Exponential Adams-Bashforth second-order timestepper.

ExponentialRungeKutta2()

Explicit second-order Runge-Kutta timestepper.

FieldWriter(path; schedule, fields=[:velocity, :temperature, :magnetic])

Schedule-driven writer that snapshots selected fields to path whenever schedule fires.

FileIC(file_path::String)

Load initial conditions from a NetCDF file.

The file is read by InitialConditions.load_initial_conditions!(field, field_type, file_path). Note that full NetCDF I/O is not yet implemented in the underlying module; at runtime a warning is printed and an analytical fallback is used instead.

field_type is inferred automatically from the field symbol passed to set_initial_condition! (:velocity, :temperature, :magnetic, or :composition).

GPU() -> GPU

Construct a GPU architecture bound to the default functional GPU backend. Errors if no GPU backend is available (CUDA extension not loaded / no device). The CUDA extension overrides _gpu_default_backend() to return a CUDABackend.

GPUPhysicalField{T,A}

Device-resident physical field: data is an (nlat, nlon, nr) array on the architecture's backend.

GPUSpectralField{T,A}

Device-resident spectral field in the DENSE (lmax+1, mmax+1, nr) (l, m, r) layout — the same layout the CPU spectral field uses and the exact form SHTnsKit's gpu_synthesis/gpu_analysis consume/produce per radial level (a dense (lmax+1, mmax+1) coefficient matrix). data_real/data_imag split the complex coefficients (mirrors the CPU container). T = real(CT) where CT is the complex type passed to allocate_gpu_spectral_field; A is the backend array type (e.g. CuArray{Float64,3} on CUDA, Array{Float64,3} on CPU).

HealthCheck(; schedule, abort=true)

Scans all simulation fields for NaN/Inf at each scheduled interval. When a non-finite value is found a warning is logged and, if abort=true, the simulation is halted by throwing.

InnerCoreAdmittance{T}

Precomputed conducting-inner-core implicit diffusion factorizations and ICB admittances, per stored harmonic degree l.

Fields

• factor: M_ic banded LU factorization per stored l.
• alpha: ICB admittance α_l per stored l.
• d1_top: dense one-sided ∂/∂r row at r=ri (row Nic of the first-derivative matrix), length Nic.
• lookup: maps a harmonic degree l to its index in factor/alpha.
• Nic: number of inner-core radial points.
• lin: the η·∇²_l operator (L) per stored l, BEFORE forming X = (1/dt)I − θL. Full-band (no boundary rows cleared); used to assemble the CNAB2 history RHS b_ic = (1/dt + (1−θ)η∇²_l) S_old.
• dt: time step used to build M_ic (must match for consistent RHS assembly).
• theta: implicit weight θ used to build M_ic.

NaNDetectionConfig

Configuration for NaN/Inf detection during simulation.

OutputSpace

Selects whether output files store physical fields, spectral fields, or the mixed layout used by GeoDynamo’s default writer.

RandomPerturbation(; amplitude, lmax, domain=nothing, seed=nothing)

Superimpose random spectral perturbations up to degree lmax ON a field. The perturbation is added to the field's existing spectral content (it does not clear the field first), so a base state applied earlier — e.g. set!(model, :temperature, AnalyticIC(:conductive)) — survives. Apply the base IC first, then the RandomPerturbation.

• amplitude – overall scale of the perturbation (required)
• lmax – maximum spherical-harmonic degree to excite (required)
• domain – DEPRECATED / unused. Accepted and forwarded for API symmetry, but the underlying randomizers currently ignore it (the off-center ball grid has no r=0 node, so no field-plane regularity zeroing is needed). Kept to avoid breaking existing call sites.
• seed – optional Int random seed for reproducibility

RungeKutta3()

Cavaglieri-Bewley/Williamson 2N-storage third-order IMEX Runge-Kutta timestepper. Nonlinear terms are advanced with the three-stage low-storage RK3 recurrence, while diffusion is solved implicitly at each substage.

SHTnsBuffers

Concrete mutable struct holding all reusable work arrays for SHTnsKit transforms. Replaces the previous Dict{Symbol,Any} buffer cache with named, type-stable fields.

Fields that are nothing at construction time are lazily allocated on first use.

SHTnsBuffers()

Construct a fully-uninitialized SHTnsBuffers with all fields set to nothing.

SHTnsPhysField{T}

A scalar field represented on the physical (θ, φ, r) grid.

Storage Layout

Field values stored in a 3D array:

• Dimension 1: Latitude index (1 to nlat), Gauss-Legendre points
• Dimension 2: Longitude index (1 to nlon), uniform spacing
• Dimension 3: Radial level index

Pencil Orientation

The pencil field specifies how data is distributed across MPI processes. Different pencil orientations are optimal for different operations:

• phi pencil: Optimal for FFTs (all longitudes local)
• theta pencil: Optimal for Legendre transforms (all latitudes local)
• r pencil: Optimal for radial operations (all radii local)

SHTnsSpecField{T}

A scalar field represented in spectral space as spherical harmonic coefficients.

Storage Layout

Coefficients are stored as separate real and imaginary parts in 3D arrays:

• Dimension 1: Spherical harmonic degree slot (l + 1)
• Dimension 2: Spherical harmonic order slot (m + 1)
• Dimension 3: Radial level index

Invalid (l,m) slots in the rectangular storage grid are ignored by the local spectral-mode map.

Boundary Conditions

Each (l,m) mode can have independent boundary conditions at the inner and outer radial boundaries. Common types:

• DIRICHLET: Specified value at boundary
• NEUMANN: Specified derivative at boundary

Fields

• config: SHTnsKit configuration (grid and transform parameters)
• nlm: Total number of spectral modes
• data_real, data_imag: Real/imaginary parts of coefficients
• pencil: PencilArrays pencil defining data distribution
• bc_type_inner/outer: Boundary condition type for each mode
• boundary_values: Boundary values [2, nlm] for inner/outer

SHTnsVectorField{T}

A vector field represented by its three spherical components in physical space.

Components:

• r_component: Radial component v_r(θ, φ, r)
• θ_component: Latitudinal component v_θ(θ, φ, r)
• φ_component: Longitudinal component v_φ(θ, φ, r)

Each component is a SHTnsPhysField, potentially with different pencil orientations for optimal computation of different operations.

SimulationProgress(; schedule)

Logs simulation time and step count at each scheduled interval.

SolenoidalMonitor(; schedule, threshold=1e-10)

Runs the solver's solenoidal-constraint diagnostic at each scheduled interval, logging the velocity (and magnetic) divergence norms and warning when the L∞ divergence exceeds threshold.

SolverBackend

Immutable description of the backend resources used by a solver run.

It records the selected architecture, SHTnsKit transform configuration, radial domains, and MPI rank/process metadata. create_solver_backend(...) is the normal constructor used by the public solver initialization path.

SolverERK2BoundarySide{T}

Boundary condition descriptor for one radial endpoint in an ERK2 solve.

type selects the physical condition, stencil supplies derivative rows for Neumann-like constraints, and the correction fields encode the l-dependent terms used by stress-free and insulating boundary formulas.

SolverERK2BoundarySpec{T}

Pair of inner/outer ERK2 boundary descriptors plus optional mode-dependent endpoint values.

Mode values are used for cases such as rotating inner-core toroidal velocity, where only selected (l,m) modes carry a nonzero endpoint value.

SolverERK2BoundarySpec{T}(inner, outer[, mode_values...])

Construct a boundary pair, inferring the concrete mode-value slot types from the arguments; with no mode-value arguments, all slots are Nothing.

SolverERK2FieldBuffers{T}

Workspace for the two-stage ERK2 update of one spectral field.

The buffers store linear propagation, first-stage nonlinear terms, provisional stage values, stage nonlinear terms, and reusable radial work vectors.

SolverERK2FieldBuffers(u, nl, cache)

Allocate ERK2 work buffers matching one spectral field, its nonlinear term, and the selected stage cache.

SolverFields{T}

Convenient grouped view of the fields that actively participate in a solver run.

This separates "allocated in the runtime" from "enabled by the current solver parameters", which keeps the timestep code simpler when optional magnetic or compositional physics are disabled.

SolverParameters

Internal solver configuration using the same names accepted by the public model constructors: nr, lmax, mmax, Ek, Pr, Pm, Sc, and related descriptive keywords.

SolverRuntime{T}

The concrete field and workspace objects stepped by the rewritten solver.

SolverState wraps this together with parameters, topography state, timestep matrices, caches, and diagnostics.

SolverState{T}

Top-level state object for the rewritten GeoDynamo solver.

It combines:

• SolverParameters
• backend and runtime objects
• active field views
• topography / Stefan state
• implicit matrices and timestep caches
• energy and solenoidal diagnostics

initialize_simulation(Float64, params) returns this type for the new solver path.

SolverTimestepState

Minimal mutable state needed by the solver timestep loop.

This tracks the current time, timestep, step counters, and simple convergence status without exposing the larger legacy runtime state.

SolverTopographyState{T}

Topography and Stefan-condition state owned by the rewritten solver.

The solver keeps this separate from the core field runtime so topographic coupling can be enabled or updated without rebuilding the spectral fields.

SolverTransformBuffers{T}

Typed scratch storage replacing TransformWorkspace{T}.cache::Dict{Symbol,Any}. Each field corresponds to a key formerly stored in the Dict.

SpecifiedTimes(times...)

Schedule that fires once when simulation time reaches (or first passes) each of the given times. Times are sorted and de-duplicated.

ThetaMethod(; theta=0.5)

Single-parameter implicit theta timestepper.

TransformWorkspace{T}

Scratch storage owned by one solver runtime for transform-side gather/scatter work.

GPU-marked runs route these allocations through backend hooks so the same solver code can use either CPU or backend-provided scratch storage.

ZeroIC()

Leave the field at its default (zero) initial state. This is a no-op: after initialize_solver_state all fields are already zero-initialized.

_apply_scalar_implicit_update!(state, 𝔽, key, diffusivity, bc_code,
                               solve_step!, etd_entry)

Field-agnostic CNAB2/EAB2/default implicit update for a scalar field. The shim supplies the field-specific cache/matrix key (:temperature / :composition), the diffusivity (Pm/Pr / Pm/Sc), the boundary bc_code, the per-field solve wrapper solve_step!, and the EAB2 cache etd_entry (state.timestep_caches.etd_temperature / etd_composition).

_apply_solver_implicit_updates_sequential!(state)

Apply every active field's implicit update on the current task.

_apply_solver_implicit_updates_threaded!(state)

Apply independent field implicit updates concurrently with Julia tasks.

This is used only for schemes whose update kernels do not require ordered MPI collectives across fields.

_build_disttranspose_plan(cfg) -> SHTnsKit.DistTransposePlan

Build (but do not cache) a DistTransposePlan for cfg. The plan is built over cfg.pencils.θ_comm with nlev = nr_local (the number of radial levels owned by this rank).

_build_disttranspose_scratch(cfg, plan) -> NamedTuple

Build the pair of scratch PencilArrays used for the Alm ↔ spec_solve transpose, keyed by cfg. Both arrays are pre-allocated once and reused. Also builds persistent Transposition plans so that transpose! never allocates MPI communication buffers on repeated calls.

Returns (; pen_alm_r, pen_solve_r, almr, solve, t_fwd, t_bwd) where:

• t_fwd : Transposition(solve, almr) — forward (almr → solve, for to_spec_solve)
• t_bwd : Transposition(almr, solve) — backward (solve → almr, for from_spec_solve!)

_callback_schedule(cb)

Returns the AbstractSchedule associated with any callback object.

_clear_p3_transform_caches!(config)

Clear all five Phase-3 scratch slots for config. All slots live on config._buffers (set to nothing directly): the DistTransposePlan, DistTranspose scratch, m-bridge, scalar-scratch, and vector-scratch. Called by clear_buffer_cache! so transient configs (e.g. across a test suite) do not accumulate stale entries.

_copy_physical2_to_spatial2!(fp1, fp2, pd1, pd2)

Copy two physical arrays to two spatial arrays in one pass (Vt+Vp vector analysis).

_copy_physical_to_spatial!(fp, pd)

Copy pd[i, j, k] → fp[j, i, k] (phys → fspatial/Vt/Vp, analysis direction). pd is (nlat, nlon, nlevpd), fp is (nlon, nlat, nlevfp).

_copy_spatial2_to_physical2!(pd1, pd2, fp1, fp2)

Copy two spatial arrays to two physical arrays in one pass (Vt+Vp vector synthesis).

_copy_spatial_to_physical!(pd, fp)

Copy fp[j, i, k] → pd[i, j, k] (fspatial/Vt/Vp → phys, synthesis direction). pd is (nlat, nlon, nlevpd), fp is (nlon, nlat, nlevfp).

_energy_diagnostics(model) -> (kinetic, magnetic)

Compute the global kinetic energy (and magnetic energy, if the model carries a magnetic field) from the current solver state's spectral coefficients. Returns a NamedTuple of Float64 energies; magnetic is 0.0 when no magnetic field is present.

_ensure_etd_cache!(caches, field, ν, T, domain)

Ensure a typed EAB2/ETD cache entry exists for one solver field.

The cache is rebuilt when diffusivity or radial resolution changes; otherwise the existing per-degree operator/factorization map is reused.

_fire_callback!(cb::Callback, sim)

Calls cb.func(sim).

_fire_callback!(cb::EnergyDiagnostics, sim)

Compute and log the current kinetic (and magnetic) energy of the model from its spectral toroidal-poloidal coefficients.

_fire_callback!(cb::HealthCheck, sim)

Scan all fields for NaN/Inf. If any are found, log a warning; when cb.abort is true, throw to halt the simulation.

_fire_callback!(cb::SimulationProgress, sim)

Logs the current simulation time and step number.

_fire_callback!(cb::SolenoidalMonitor, sim)

Run the solver's solenoidal-constraint diagnostic and log the divergence norms, warning when the velocity (or magnetic) L∞ divergence exceeds cb.threshold.

_get_disttranspose_scratch(cfg, plan)

Return the cached scratch NamedTuple for cfg, building it on the first call. The scratch is stored in cfg._buffers.disttranspose_scratch (a mutable field on the config-owned SHTnsBuffers). Subsequent calls are O(1) field reads.

_get_influence_cache(domain::RadialDomain, ∂r)

Get or create cached influence functions and matrix for given domain. Note: ∂r should be a BandedMatrix from numerics/banded_operators.jl.

_get_or_build_erk2_cache(existing, label, diffusivity, T, config, domain, dt; ...)

Build or reuse an ERK2 stage cache for velocity-like fields.

Callers own the storage location; this helper only decides whether the existing cache still matches the current grid, timestep, diffusivity, and method flags.

_get_or_build_erk2_influence_entry(existing, T, config, domain, diffusivity, dt, velocity_bc_code; theta)

Build or reuse the velocity-poloidal influence correction cache.

The cache key includes a hash of the radial grid so boundary-correction operators are refreshed when the domain geometry changes.

_get_or_build_erk2_scalar_cache(existing, label, diffusivity, T, config, domain, dt, boundary_condition; ...)

Build or reuse an ERK2 stage cache for scalar fields.

_get_or_build_velocity_workspace!(𝒰, nr; nthreads) -> VelocityWorkspace

Return the velocity field's own scratch workspace, building (and caching on the field) a fresh zeroed one when absent or sized for fewer threads / a different radial resolution. The workspace is owned by 𝒰 — never a module global — so scratch is never shared between solver states. Cross-state sharing previously let one state's leftover values leak into another's radial-profile kernels, which made the lmax=2 magnetic-poloidal solve nondeterministic under the full suite (finite alone, occasionally non-finite in CI).

_get_tau_cache(domain::RadialDomain)

Get or create cached tau polynomials and derivatives for given domain.

_health_check(model) -> (has_issue, fields)

Scan every spectral field of the model (velocity, temperature, and magnetic / composition when present) for NaN/Inf values. Returns a NamedTuple with has_issue::Bool and fields::Vector{String} naming the offending components.

_ic_build_bic(a::InnerCoreAdmittance, l, S_old) -> Vector

Assemble the CNAB2 history RHS for the inner-core diffusion solve at degree l:

b = (1/dt) S_old + (1−θ) · η∇²_l · S_old

using the SAME η∇²_l (a.lin), dt, and θ that built M_ic. Boundary rows (b[1], b[Nic]) are left as produced by the full-band Laplacian and are meant to be overridden by the caller (regularity at r=0, prescribed ICB value).

_load_restart_file(filepath, tracker, config; pencils=nothing)

Load simulation state from a specific restart NetCDF file path using parallel I/O. All ranks open the file collectively and read their local slices.

_magnetic_conducting_history_flux(magnetic, ic_spec, adm) -> (φ0_real, φ0_imag)

Build the conducting-inner-core ICB history flux φ0 for every local magnetic mode (l,m) of ic_spec (the inner-core scalar toroidal_ic or poloidal_ic), using its OLD radial profile over inner-core indices 1..Nic.

φ0_l = inner_core_history_flux(adm, l, S_ic_old) is the ICB radial-derivative contribution from the inner-core CNAB2 history with a zero ICB value; paired with the admittance α_l baked into the outer-core Robin inner row it enforces derivative continuity across the ICB. Returns mode-indexed real/imag vectors (length nlm) that are MPI all-reduced so every rank feeds the same boundary values to the matrix rows. l=0 is skipped (magnetic has no l=0 mode).

_magnetic_conducting_reconstruct!(oc_spec, ic_spec, adm)

After the outer-core solve, reconstruct the inner-core radial profile for every local magnetic mode and write it into ic_spec (toroidal_ic or poloidal_ic).

For each mode the ICB value g is the outer-core solution oc_spec at radial index 1 (the outer-core inner point coincides with the ICB), and S_ic_new = reconstruct_inner_core(adm, l, g, S_ic_old) is solved with regularity S(0)=0 and ICB Dirichlet value S(ri)=g. Value continuity across the ICB holds by construction (g is shared); the result is written over inner-core radial indices 1..Nic. l=0 is skipped.

_populate_radial_operators!(domain) -> domain

Fill domain.dr_matrices and domain.radial_laplacian with the finite-difference operators for this grid, in place.

The RadialDomain constructors allocate these as zero matrices. Without this step they stay zero: every current consumer recomputes via create_derivative_matrix/create_radial_laplacian, but any code that reads the stored fields directly would silently get zeros (wrong derivatives, no error). Populating them keeps the cached operators consistent with a fresh recompute.

Degenerate grids with N < length(dr_matrices) + 1 lack enough stencil points for the highest-order derivative; their operators are left zero (as before). Physical runs use N far larger than the stencil order, so this only affects toy domains that never build radial operators anyway.

_prepare_solver_eab2_caches!(state)

Ensure all active ExponentialAdamsBashforth2 exponential-action caches match the current parameters.

_run_callbacks!(sim)

Iterates over sim.callbacks, builds a _ScheduleContext from the current simulation state, and fires each callback whose schedule returns true from should_fire.

_run_output_writer!(ow::CheckpointWriter, sim, ctx)

Fires a checkpoint write if the writer's schedule fires for ctx.

NOTE: The underlying write_restart! in io/writer.jl requires a Dict{String,Any} of fields (from extract_all_fields), a TimeTracker, and an OutputConfig. When MPI is not initialized this falls back to logging.

_run_output_writer!(ow::FieldWriter, sim, ctx)

Fires a field snapshot write if the writer's schedule fires for ctx.

NOTE: The underlying write_fields! in io/writer.jl requires a TimeTracker and OutputConfig and must be called collectively by all MPI ranks. When MPI is not initialized (e.g. unit tests) this falls back to logging an info message without performing I/O.

_run_output_writers!(sim)

Iterates over sim.output_writers, builds a _ScheduleContext from the current simulation state, and fires each writer whose schedule returns true.

_select_output_fields(all_fields, selectors)

Return the subset of all_fields (the dict from extract_all_fields) selected by selectors (e.g. [:velocity, :temperature]). An empty selector list means "write everything". Unknown selectors and keys absent from all_fields are skipped silently.

_shtns_make_transpose(pair)

Create a PencilArrays transpose plan between two pencil configurations. Used internally to set up efficient data redistribution operations.

Arguments

• pair: A Pair of source and destination Pencil objects (src => dest)

Returns

• A Transposition object that can be used with mul! for data redistribution

_solenoidal_diagnostics(model) -> (velocity_l2, velocity_linf, magnetic_l2, magnetic_linf)

Run the solver's solenoidal-constraint diagnostic against the current state and return the freshly recorded divergence norms for the velocity (and magnetic) fields. This records one sample into model.state.solenoidal_monitor.

NOTE: the divergence norms come from compute_divergence_spectral, which is currently a deferred solver stub returning zero; the toroidal-poloidal representation is solenoidal by construction so this metric is ~0 until that primitive is implemented. The callback itself is fully wired: it runs the real diagnostic pipeline and threshold-checks the recorded values.

_solver_can_thread_implicit_updates(timestepper)

Return whether field implicit updates can run on Julia threads.

Field update kernels reach MPI collectives (ExponentialAdamsBashforth2 directly; others via influence and transpose paths). With more than one rank, issuing those collectives from multiple Julia threads lets the per-rank ordering diverge and the collectives mismatch, which deadlocks. So restrict threaded field updates to single-rank runs regardless of scheme; multi-rank runs use the sequential path.

_solver_compute_scalar_nonlinear!(𝔽, vel_fields, outer_core_domain, ws;
                                  add_internal_sources)

Field-agnostic scalar nonlinear assembly: zero work arrays, compute gradients, transform field+gradients to physical space, form advection (and, when add_internal_sources, the radial internal-source profile), then transform the nonlinear term back to spectral. Both temperature and composition pass add_internal_sources = true so their respective source fields (internal_heating / compositional_source) enter the evolution.

_solver_solve_scalar_implicit_step!(solution, rhs, matrices; bc_inner=nothing, bc_outer=nothing, ...)

Solve one scalar-field implicit radial system for each local spectral mode.

Boundary vectors are mode-indexed global arrays. Missing boundary vectors default to homogeneous values for both real and imaginary parts.

_solver_uses_two_step_history(timestepper)

Return true for timestep schemes that need previous nonlinear terms.

_spectral_curl_torpol!(dst_tor_r, dst_tor_i, dst_pol_r, dst_pol_i,
                       src_tor_r, src_tor_i, src_pol_r, src_pol_i,
                       l_factors, d1_matrix, d²_matrix, domain, config, T)

Shared spectral curl for toroidal-poloidal fields: (∇×V)tor = [l(l+1)/r² - d²/dr² - 2/r d/dr] Vpol (∇×V)pol = -l(l+1)/r² Vtor

Used by both current density (j = ∇×B) and induction curl (∇×(u×B)).

_sync_solver_history!(previous, current)

Copy spectral field storage from current into the matching previous buffer.

This helper works on the real and imaginary parent arrays so it covers both local PencilArray views and ordinary arrays.

_sync_solver_nonlinear_histories!(state, magnetic_enabled)

Copy all active nonlinear terms into their previous-step buffers.

Temperature and velocity are always present; magnetic and composition histories are updated only when those fields are enabled in the solver state.

add_internal_sources_local!(𝔽::AbstractScalarField{T}, domain::RadialDomain) where T

Add volumetric sources (completely local operation). This works for any scalar field with radial source profile.

allocate_gpu_physical_field(T, arch, config, nr) -> GPUPhysicalField

Allocate a zero-filled (nlat, nlon, nr) physical field on arch's backend.

allocate_gpu_spectral_field(CT, arch, config, nr) -> GPUSpectralField

Allocate a zero-filled dense (lmax+1, mmax+1, nr) split-complex spectral field on arch's backend. CT is the complex element type (ComplexF64); storage is its real part type (Float64).

analyze_load_balance(pencil::Pencil)

Analyze and report load balance for a given pencil decomposition.

apply_banded_full!(out, B, v)

Apply a banded radial operator to a full vector without materializing the dense matrix.

This is the single implementation used by both the solver internals and the low-level public API.

apply_flux_bc_direct!(spec_real, spec_imag, slot, lm_idx,
                     𝔽::AbstractScalarField, domain, r_range)

Direct modification of boundary values to approximately satisfy flux BC. This is the simplest but least accurate method.

apply_flux_bc_influence_matrix!(spec_real, spec_imag, lm_idx,
                               apply_inner, apply_outer, 𝔽::AbstractScalarField,
                               domain, r_range, owns_mode::Bool)

Apply flux boundary conditions using the influence matrix method.

MPI Safety

The owns_mode parameter indicates whether this process owns the lm mode. All processes must call this function for each mode to ensure Allreduce is called the same number of times by all processes (prevents deadlock).

apply_flux_bc_tau!(spec_real, spec_imag, lm_idx,
                   apply_inner, apply_outer, 𝔽::AbstractScalarField,
                   domain, r_range, owns_mode::Bool)

Apply flux boundary conditions using the tau method. This is the generalized version that works with any scalar field.

MPI Safety

The owns_mode parameter indicates whether this process owns the lm mode. All processes must call this function for each mode to ensure Allreduce is called the same number of times by all processes (prevents deadlock).

apply_geometric_factors_spectral!(ws::GradientWorkspace{T}, 𝔽::AbstractScalarField{T}, domain::RadialDomain) where T

Apply geometric factors (1/r, 1/(r sin θ)) in spectral space. For gradients in spherical coordinates. This is generic.

GeoDynamo.apply_influence_matrix_correction!(result, influence, bc_inner_val=0, bc_outer_val=0)

Public wrapper for applying a single radial-profile influence correction.

apply_scalar_boundary_parameters!(field, boundary_conditions)

Install scalar boundary values from SolverParameters into the field-owned boundary storage used by CNAB2, EAB2, and ERK2.

Parameter-specified scalar BCs are spatially uniform, so only the (l,m)=(0,0) mode receives a nonzero endpoint value. File-based spectral BCs can later replace this with per-mode real and imaginary values.

The transform is orthonormal (Y_0^0 = 1/√(4π)), so a uniform physical value v maps to the (0,0) spectral coefficient v·√(4π); the boundary endpoints are scaled accordingly (same convention as bcs/topography/topography_data.jl).

apply_scalar_flux_bc_spectral!(𝔽::AbstractScalarField{T}, domain::RadialDomain;
                               method::Symbol=:tau) where T

Apply flux boundary conditions to a scalar field in spectral space. Methods available: :tau (most robust), :influence_matrix, :direct (simplest).

MPI Safety

Uses global loop bounds (1:nlm) to ensure all processes call MPI collectives the same number of times, preventing deadlock with uneven lm distribution.

apply_solver_erk2_stage!(buffers, u)

Overwrite u with the provisional half-step ERK2 stage stored in buffers.

apply_solver_implicit_step!(state)

Run one configured implicit/IMEX timestep update for all active fields.

The dispatcher handles first-step history bootstrap, prepares ExponentialAdamsBashforth2 caches when needed, delegates ExponentialRungeKutta2 to the staged integrator, and rolls nonlinear histories after a successful update.

apply_solver_influence_matrix_correction!(result, influence, bc_inner_val=0, bc_outer_val=0)

Project one dense radial profile back onto the requested endpoint constraints using a precomputed two-column influence operator.

apply_solver_parameters!(params)

Record params as the active solver parameter set used by shared helpers.

apply_solver_velocity_poloidal_influence_correction!(field, influence_matrices, config)

Apply the velocity-poloidal endpoint influence correction to all local spectral modes in a distributed field.

apply_tau_correction!(profile::Vector{T}, tau_coeffs, domain::RadialDomain) where T

Add tau correction to the radial profile.

GeoDynamo.apply_velocity_poloidal_influence_correction!(field, influence_matrices, config)

Public wrapper for applying influence corrections to velocity-poloidal fields.

arch_of(a) -> AbstractArchitecture

The GeoDynamo architecture matching an array's compute backend. Dispatches on the KernelAbstractions backend, so any CPU-backed array (Array, SubArray, …) maps to CPU() and a device array (e.g. CuArray) to GPU(backend) — robust to wrapped/viewed arrays (unlike an isa Array test).

arch_zeros(arch, FT, dims...)

Allocate a zero-filled array on the given architecture.

banded_to_dense(matrix)

Materialize a BandedMatrix as a dense matrix.

This is mainly a debugging/compatibility helper; production timestep solves prefer banded operators and factorizations.

batch_magnetic_transforms!(ℬ::SHTnsMagneticFields{T}) where T

Perform batched transforms for better cache efficiency using transforms/spectral.jl.

batch_spectral_to_physical!(specs, physs)

Compatibility wrapper that calls batch_shtnskit_transforms! for batched spectral→physical transforms using SHTnsKit with PencilArrays/MPI.

timestepper isa RungeKutta3 && return integrate_solver_cb3_step!(state)

bootstrap_solver_history!(state, timestepper, magnetic_enabled)

Seed previous nonlinear histories for two-step schemes on their first step.

CNAB2, ExponentialAdamsBashforth2, and ExponentialRungeKutta2 all need a previous nonlinear term. At startup there is no true previous step, so the current nonlinear term is copied once to avoid a special-case branch in each update kernel.

build_gpu_erk2_state(cpu_state) -> NamedTuple

Assemble the ERK2 device-state pack from a CPU SolverState configured with the ERK2 timestepper. Builds (or reuses) the same memoized propagator caches and boundary specs the CPU integrate_solver_erk2_step! uses, then packs them into dense batched arrays for gpu_erk2_solver_step!. Combine with the base bundle from build_gpu_solver_state:

st  = ...                      # CPU SolverState, ERK2 timestepper
solver_step!(st)               # warm-up (physical buffers, caches)
gst = build_gpu_solver_state(st)
erk = build_gpu_erk2_state(st)
gpu_erk2_solver_step!(gst, erk)

Arrays are host Arrays; move with gpu_to_device alongside the base bundle. Boundary endpoint VALUES are baked at pack time (time-constant BCs only).

build_gpu_solver_state(cpu_state) -> NamedTuple

Assemble the gpu_solver_step! device-state bundle from a CPU SolverState (insulating magnetic, homogeneous BCs). Arrays are on the CPU (Array) backend; move to a device with on_architecture(GPU(), …) for a GPU run. The returned NamedTuple matches the layout consumed by gpu_solver_step!: per-field tor/pol/scalar bundles (dense spectral + prevnl + lin/lu + BC rows), shared operators (`nlopsvel/nlopsmag,d1/mvals/rinv/rvec), coupling factors, the velocity-poloidal influence pack, and the persistent lagged physical buffers (Tphys/Cphys/B*/J*`).

The CPU spectral storage is slot-packed; cpu_spectral_to_dense scatters it to the dense (lmax+1, mmax+1, nr) layout the GPU kernels use. The physical lag buffers are read from the CPU physical fields AS THEY STAND — to reproduce the one-step velocity lag, build the device state AFTER one warm-up solver_step!.

build_influence_matrix(G_inner, G_outer, ∂r, domain)

Construct a 2×2 matrix mapping influence amplitudes to boundary flux errors. Rows correspond to (inner, outer) boundaries; columns to (inner, outer) influence functions.

build_magnetic_implicit_matrices_conducting(T, cfg, domain, ic_domain, dt; theta)

Build the magnetic toroidal/poloidal implicit matrices for a conducting inner core together with the inner-core ICB admittances.

For each component the inner-core diffusion operator and its ICB admittance α_l are precomputed (create_inner_core_admittance), and the outer-core matrices are built with the conducting Robin inner row (∂/∂r − α_l) S = φ0 via the inner_alpha kwarg. The magnetic diffusivity (1.0), dt, and theta MUST match the values used to shift the outer-core system matrices so the CNAB2 history RHS is consistent across the ICB.

Returns (tor, pol, admittance) where admittance is a NamedTuple{(:tor,:pol)} of InnerCoreAdmittance{T} objects.

build_rhs_cnab2!(rhs, un, nl, nl_prev, dt, matrices)

Build RHS for CNAB2 IMEX: rhs = un/dt + (1-θ)·L·un + (3/2)·nl − (1/2)·nl_prev, where θ = matrices.theta and L is the diffusivity-scaled linear operator.

MPI Safety

Uses global loop bounds (1:nlm) to ensure all processes call Allreduce the same number of times, preventing deadlock with uneven lm distribution.

build_solver_erk2_magnetic_pol_bc(T, domain; inner_regularity=false)

Create insulating boundary descriptors for magnetic poloidal fields. With inner_regularity = true (ball / full-sphere geometry) the inner side is the center-regularity row P′(r₁) = (l+1)·P(r₁)/r₁ (P ~ r^{l+1}); the insulating outer side is unchanged.

build_solver_erk2_magnetic_tor_bc(T, domain; inner_regularity=false)

Create boundary descriptors for magnetic toroidal fields: homogeneous Dirichlet on both sides (insulating shell walls). With inner_regularity = true (ball / full-sphere geometry) the inner side is the center-regularity row t′(r₁) = l·t(r₁)/r₁; the outer Dirichlet is unchanged.

build_solver_erk2_scalar_bc(T, domain, boundary_condition; inner_regularity=false)

Translate scalar boundary-condition codes into an ERK2 boundary specification.

Boundary codes follow the existing scalar convention: DD, DN, ND, and NN for codes 1 through 4. With inner_regularity = true (ball / full-sphere geometry) the inner side is the center-regularity row Θ′(r₁) = l·Θ(r₁)/r₁ regardless of the code; the outer side is unchanged.

build_solver_erk2_velocity_tor_bc(T, domain, velocity_bc_code; config=nothing, rot_omega=0.0, inner_regularity=false)

Create ERK2 boundary descriptors for the velocity toroidal component.

When a rotating inner core is requested, the (l=1, m=0) mode gets a mode-dependent inner boundary value. With inner_regularity = true (ball / full-sphere geometry) the inner side is the center-regularity row t′(r₁) = l·t(r₁)/r₁ regardless of the code, and the rotating-inner-core mode values are skipped (there is no inner core in a ball).

build_∂θ(::Type{T}, config::SHTnsKitConfig) where T

Build sparse matrix for θ-derivatives in spectral space. This matrix couples different l modes with the same m.

check_field_for_nan(field_data, field_name, config, step)

Check a field for NaN or Inf values. Returns (has_nan, has_inf, nan_count, inf_count).

check_parallel_netcdf_support(comm)

Verify that parallel NetCDF (MPI-IO via HDF5) is available at runtime. Called once at initialization. Errors immediately if not supported.

check_simulation_state_for_nan(state, step; config=DEFAULT_NAN_CONFIG)

Comprehensive NaN/Inf check across all simulation fields. Returns true if any NaN/Inf is detected.

check_solver_health!(state)

Run periodic health checks on the runtime state.

Currently this checks for NaNs every ten solver steps.

check_spectral_field_for_nan(field, field_name, config, step)

Check both real and imaginary parts of a spectral field.

cleanup_old_files(output_dir, keep_last_n=10)

Remove older history files from output_dir, keeping the most recent keep_last_n files by history number.

compat_normalize_old_erk2_cache_entry(entry)

Extract an ERK2StageCache from legacy cache bundle entries.

Older code sometimes stores cache entries directly and sometimes under a :cache key; unsupported entries return nothing.

compute_all_gradients_spectral!(𝔽::AbstractScalarField{T}, domain::RadialDomain, ws::GradientWorkspace{T}) where T

Compute all gradient components (θ, φ, r) in spectral space. This is the main driver function that works for any scalar field.

compute_all_nonlinear_terms!(𝒰, temp_field, comp_field, mag_field, domain)

Compute all nonlinear forcing terms for the momentum equation.

Governing Equation (Dimensional)

In a rotating reference frame with angular velocity Ω:

∂u/∂t = -u×ω - ∇p/ρ + 2Ω(ẑ×u) + ν∇²u + (αg/ρ)T·r̂ + (Δρg/ρ)C·r̂ + (1/μ₀ρ)(∇×B)×B

where:

• u: velocity, ζ = ∇×u: vorticity
• Ω: rotation rate, ẑ: rotation axis
• ν: kinematic viscosity
• α: thermal expansion coefficient, g: gravity
• T: temperature perturbation, C: composition perturbation
• B: magnetic field, μ₀: permeability

Non-Dimensionalization (Magnetic Diffusion Time Scaling)

Length scale: d (shell thickness or ball radius) Time scale: τ = d²/η (magnetic diffusion time) Velocity scale: U = η/d Magnetic field scale: B₀ Temperature scale: ΔT

Dimensionless parameters:

• E = ν/(Ω d²): Ekman number
• Pm = ν/η: Magnetic Prandtl number
• Pr = ν/κ: Prandtl number
• Sc = ν/D: Schmidt number
• Ra = (αgΔT d³)/(νκ): Rayleigh number
• Ra_C = (Δρg d³)/(ρνD): Compositional Rayleigh number

Non-Dimensional Momentum Equation (magnetic diffusion time scaling)

E·∂ũ/∂τ = E(ũ×ω̃) - ẑ×ũ - ∇p̃* + (Pm/Pr)·Ra·T̃·r·r̂ + (Pm/Sc)·Ra_C·C̃·r·r̂ + (1/Pm)(∇×B̃)×B̃ + E∇²ũ

where:

• τ = L²/η (magnetic diffusion time scaling)
• E = ν/(2ΩL²) is the Ekman number
• r factor in buoyancy represents linear gravity profile g(r) ∝ r

Implementation Notes

The explicit RHS entering the time integrator is: RHS = E·(ũ×ω̃) - (ẑ×ũ) + (Pm/Pr)·Ra·T̃·r·r̂ + (Pm/Sc)·Ra_C·C̃·r·r̂ + (1/Pm)(∇×B̃)×B̃

Coefficients (magnetic diffusion time scaling, mass coeff = E on du/dt):

• Advection: E
• Coriolis: 1 (no scaling)
• Thermal buoyancy: (Pm/Pr) * Ra * r (with radial factor)
• Compositional buoyancy: (Pm/Sc) * Ra_C * r (with radial factor)
• Lorentz: 1/Pm

Viscous diffusion is treated implicitly with coefficient E (Ekman number). Mass coefficient E is applied in the time-stepping matrices.

compute_boundary_fluxes(profile::Vector{T}, ∂r, domain::RadialDomain) where T

Compute flux (dT/dr) at both boundaries using the derivative matrix. Note: ∂r should be a BandedMatrix{T} from numerics/banded_operators.jl.

compute_chebyshev_polynomial(n::Int, domain::RadialDomain)

Compute Chebyshev polynomial T_n on the radial grid.

compute_composition_nonlinear!(𝔽::SHTnsCompositionField{T},
                               vel_fields, outer_core_domain::RadialDomain,
                               ws::GradientWorkspace{T};
                               geometry::Symbol) where T

Compute the nonlinear advection term -u·∇C for the composition equation.

This function implements the EXPLICIT part of the composition equation: ∂C/∂t + u·∇C = (Pm/Sc) ∇²C

The advection term -u·∇C is computed using the following optimized workflow:

1. Compute ∇C in spectral space (no MPI communication required)
2. Batched transform of C and ∇C to physical space
3. Compute advection -u·∇C locally in physical space
4. Transform result back to spectral space
5. Apply boundary conditions

Arguments

• 𝔽: Composition field structure to update
• vel_fields: Velocity field structure (provides u for advection)
• outer_core_domain: Radial domain information
• geometry: Geometry type (:shell or :ball) for transform selection

Side Effects

• Updates 𝔽.nonlinear with the computed advection term
• Updates 𝔽.gradient with ∇C in physical space
• Applies boundary conditions to spectral representation

Performance Notes

• Gradients computed in spectral space avoid MPI communication
• Batched transforms minimize data transfer overhead
• Physical space operations are embarrassingly parallel

compute_diagnostics(fields, field_info)

Compute scalar diagnostics for available physical and spectral fields.

Physical fields get global mean, standard deviation, minimum, and maximum. Spectral fields get globally reduced energy, RMS, and maximum coefficient magnitude, with extra degree-wise diagnostics when SHTns metadata is available.

compute_divergence_spectral(tor_spec, pol_spec, domain) -> (l2, linf)

Real divergence diagnostic: synthesizes the vector field from its (T, P) potentials, re-analyzes it into QST scalars (Q from the radial component, S/T from the tangential sphtor analysis), and evaluates the per-mode spectral divergence

div_lm(r) = (1/r²)·∂_r(r²·Q_lm) − l(l+1)·S_lm/r

returning RMS (L2) and L∞ norms over modes and interior radial nodes. Going through physical space and back makes this a genuine check of the transform chain (a (T,P) pair synthesizes solenoidally by construction under the Stage-2 convention, so any nonzero divergence here is a chain defect). The spectral formula avoids the grid-space aliasing that scalar analysis of tangential components (1/sinθ structure) suffers beyond lmax for m>0 modes.

Replaces a stub that hardcoded (0.0, 0.0). Allocation-heavy (builds scratch fields per call) — diagnostic path only.

The per-mode accumulators are reduced across the communicator (SUM for the L2 numerator and count, MAX for L∞) so the returned norms are global on multi-rank runs; under a single rank / no MPI the reduction is the identity.

compute_field_energy(field_data::Array{T,3}) where T

Compute L2 energy of a 3D scalar field.

compute_influence_functions_flux(outer_core_domain::RadialDomain)

Compute influence functions for flux BCs. These are solutions to the homogeneous equation with specific BCs.

compute_inner_tau_function(domain::RadialDomain)

Compute tau function for inner boundary only (zero derivative at outer).

compute_magnetic_helicity(ℬ::SHTnsMagneticFields{T}) where T

Compute magnetic helicity using enhanced spectral integration

compute_outer_tau_function(domain::RadialDomain)

Compute tau function for outer boundary only (zero derivative at inner).

GeoDynamo.compute_phi1_function(A, expA)

Public compatibility wrapper for the solver-local phi1 matrix-function helper.

GeoDynamo.compute_phi2_function(A, expA; l=0)

Public compatibility wrapper for the solver-local phi2 matrix-function helper.

compute_phi_gradient_spectral!(𝔽::AbstractScalarField{T}, ws::GradientWorkspace{T}) where T

Compute ∂field/∂φ using spherical harmonic properties (local operation). This is generic and works for any scalar field.

compute_radial_gradient_spectral!(𝔽::AbstractScalarField{T}, domain::RadialDomain, ws::GradientWorkspace{T}) where T

Compute ∂field/∂r using banded matrix derivative operator in spectral space. This uses the pre-computed derivative matrices from the field for optimal accuracy and efficiency.

Note: Radial data is always local (not MPI distributed). Only spectral (lm) modes are distributed across MPI processes.

compute_scalar_advection_local!(𝔽::AbstractScalarField{T}, vel_fields) where T

Compute -u·∇field in physical space (completely local operation). This works for any scalar field.

compute_scalar_energy(𝔽::AbstractScalarField{T}, outer_core_domain::RadialDomain) where T

Compute energy ∫ field² dV

compute_scalar_rms(𝔽::AbstractScalarField{T}, outer_core_domain::RadialDomain) where T

Compute RMS value of scalar field.

compute_spectral_energy_diagnostics!(diagnostics, component, real_part, imag_part, field_info)

Add degree-wise spectral energy summaries for one complex spectral component.

Requires field_info.config so coefficient rows can be mapped to spherical harmonic degree l. The function mutates diagnostics in place.

compute_tau_coefficients_both(flux_error_inner::T, flux_error_outer::T, domain::RadialDomain) where T

Compute tau polynomial coefficients for both boundaries. Uses highest two Chebyshev modes as tau functions.

compute_tau_coefficients_inner(flux_error::T, domain::RadialDomain) where T

Compute tau coefficient for inner boundary only. Uses a single tau function that doesn't affect outer boundary.

compute_tau_coefficients_outer(flux_error::T, domain::RadialDomain) where T

Compute tau coefficient for outer boundary only. Uses a single tau function that doesn't affect inner boundary.

compute_theta_gradient_spectral!(𝔽::AbstractScalarField{T}, ws::GradientWorkspace{T}) where T

Compute ∂field/∂θ using spherical harmonic recurrence relations (local operation). This is generic and works for any scalar field.

compute_theta_recurrence_coefficients(::Type{T}, config::SHTnsKitConfig) where T

Pre-compute recurrence coefficients for θ-derivatives. Store as [nlm, 2] matrix: [:, 1] for l-1 coupling, [:, 2] for l+1 coupling

compute_vector_energy(v_r, v_theta, v_phi)

Compute kinetic or magnetic energy from three physical-space vector components.

conductive_profile_solve(; domain, bc_code, inner_value, outer_value,
                         source, inner_regularity) -> Vector{Float64}

Steady l=0 (0,0)-coefficient profile c(r) solving ∇²₀ c = −source with the same boundary rows the implicit scalar matrices use (⇒ the IC is a discrete equilibrium). inner_value/outer_value are the boundary coefficients (Dirichlet value, or Neumann flux). source is a length-N vector (RHS of ∇²c = −S). The function is linear and unit-agnostic; callers pre-scale by √(4π) as needed. inner_regularity=true selects the ball-centre row (Θ′(r₁)=0 for l=0).

cpu_bc_to_dense(bc_vec, config, ::Type{T}) -> dense

Scatter a CPU per-mode (length nlm) boundary vector into a dense (lmax+1, mmax+1) array (mode (l,m) → (l+1, m+1)). nothing → all zeros.

cpu_spectral_to_dense(field_spec, config, nr, ::Type{T}) -> (dense_r, dense_i)

Scatter the CPU slot-packed spectral storage of field_spec into dense (lmax+1, mmax+1, nr) real/imag arrays (mode (l,m) → slot (l+1, m+1)). Modes absent from the local storage map are left at zero.

create_composition_matrices(config, domain, diffusivity, dt;
                             composition_bc_code, theta, T, inner_regularity)

Create implicit time-stepping matrices for the composition equation with Dirichlet/Neumann boundary rows embedded in the radial system matrix.

Pass inner_regularity = true for ball (full-sphere) geometry to impose the centre regularity condition Θ′(r₁) = l·Θ(r₁)/r₁ on the inner boundary row.

create_computation_pencils(topology, dims, config, θ_comm, r_comm, θ_ranks)

Create specialized pencils for different stages of computation.

Phase 2 (r×θ): pencil_r uses decomp (1,3) so θ(dim1) is distributed over θranks and r(dim3) is distributed over rranks, leaving φ(dim2) LOCAL. theta_phys is a 2D (nlat,nlon) pencil on the θ-subcommunicator so its θ-split matches pencil_r's θ-split within each r-group. θ_comm and r_comm are forwarded in the returned NamedTuple for sub-collective operations.

GeoDynamo.create_dirichlet_bc(T, nr, value=zero(T))

Create a public fixed-value ERK2 endpoint descriptor.

create_enhanced_metadata(state::SolverState, time, step)

Create lightweight metadata describing a solver snapshot for outputs and restart-style tooling.

GeoDynamo.create_erk2_cache(T, config, domain, diffusivity, dt; ...)

Public wrapper for constructing a generic ERK2 stage cache.

GeoDynamo.create_erk2_cache_composition(T, config, domain, diffusivity, dt, composition_bcs; ...)

Create the ERK2 cache used by composition fields.

GeoDynamo.create_erk2_cache_magnetic_poloidal(T, config, domain, diffusivity, dt; ...)

Create the ERK2 cache used by magnetic poloidal fields.

GeoDynamo.create_erk2_cache_magnetic_toroidal(T, config, domain, diffusivity, dt; ...)

Create the ERK2 cache used by magnetic toroidal fields.

GeoDynamo.create_erk2_cache_scalar(T, config, domain, diffusivity, dt, boundary_condition; ...)

Public wrapper for constructing scalar-field ERK2 caches with embedded boundary conditions.

GeoDynamo.create_erk2_cache_temperature(T, config, domain, diffusivity, dt, temperature_bcs; ...)

Create the ERK2 cache used by temperature fields.

create_erk2_config(; lmax, mmax, nlat, nlon, optimize_for_erk2=true)

Create an SHTnsKit configuration for ERK2 timestepping.

create_inner_core_admittance(T, l_values, ic_domain, diffusivity, dt; theta=0.5)
    -> InnerCoreAdmittance{T}

Build the conducting-inner-core implicit diffusion operator M_ic = (1/dt)I − θ·diffusivity·∇²_l and its ICB admittance α_l for every positive harmonic degree in l_values.

ic_domain is the inner-core ball domain (index 1 = r=0, index Nic = r=ri). l=0 is skipped (magnetic has no l=0 mode).

create_inner_core_spectral_pencil(config, reference_spec_pencil, nr_inner) -> Pencil

Spectral pencil for inner-core fields. Reuses the (l, m) process topology and decomposition of the outer-core spectral pencil but sizes the (local) radial dimension to the inner-core grid (nr_inner) instead of the outer-core nr, so toroidal_ic/poloidal_ic own exactly their physical radial extent rather than padding to nr. The mode-slot (l, m) layout is identical to the outer-core spectral pencil.

GeoDynamo.create_insulating_inner_bc(T, d1_row, r_inv)

Create a public inner-boundary insulating magnetic-poloidal descriptor.

GeoDynamo.create_insulating_outer_bc(T, d1_row, r_inv)

Create a public outer-boundary insulating magnetic-poloidal descriptor.

create_magnetic_poloidal_matrices(config, domain, diffusivity, dt;
                                  theta, T, inner_alpha, inner_regularity)

Create implicit time-stepping matrices for the poloidal magnetic component with insulating boundary conditions embedded in the matrix rows (matching Fortran magbcPol).

Boundary conditions (l-dependent, vacuum matching under B_r = λP/r²):

• Inner: (∂/∂r - (l+1)/r) BP = 0 (interior vacuum solution P ∝ r^{l+1})
• Outer: (∂/∂r + l/r) BP = 0 (exterior vacuum decay P ∝ r^{-l}; verified by the full-sphere dipole free-decay rate σ = π², test/ballbesseldecay.jl)

When inner_regularity = true (ball / full-sphere geometry) the inner row instead imposes the center regularity condition P′(r₁) = (l+1)·P(r₁)/r₁ (β = l+1, P ~ r^{l+1}); the outer row is unchanged.

create_magnetic_toroidal_matrices(config, domain, diffusivity, dt;
                                  theta, T, inner_alpha, inner_regularity)

Create implicit time-stepping matrices for the toroidal magnetic component with insulating boundary conditions embedded in the matrix rows (matching Fortran magbcTor).

Both boundary rows use identity (BT = 0) for insulating exterior/interior. When inner_regularity = true (ball / full-sphere geometry) the inner row instead imposes the center regularity condition t′(r₁) = l·t(r₁)/r₁ (β = l); the outer row is unchanged.

GeoDynamo.create_neumann_bc(T, d1_row, value=zero(T); l0_dirichlet=false)

Create a public first-derivative ERK2 endpoint descriptor.

GeoDynamo.create_noslip_pol_bc(T, d1_row)

Create a public poloidal-velocity no-slip endpoint descriptor.

create_parallel_netcdf(filename, config, field_info, metadata, comm)

Open a new NetCDF file in parallel mode. All ranks call this collectively. Defines global dimensions and variables based on field_info.

create_pencil_decomposition_shtnskit(nlat, nlon, nr, sht_config, comm, optimize; lmax, mmax)

Create PencilArrays decomposition optimized for spherical harmonic transforms.

The Pencil Decomposition Strategy

In a pencil decomposition, 3D data is distributed across MPI processes such that one dimension is always fully local (not split across processes). This enables efficient operations along that dimension without MPI communication.

Physical Space Pencils (nlat × nlon × nr)

theta pencil:  [θ local] × [φ distributed] × [r distributed]
               → Best for operations along latitude (Legendre transforms)

phi pencil:    [θ distributed] × [φ local] × [r distributed]
               → Best for FFTs along longitude

r pencil:      [θ distributed] × [φ distributed] × [r local]
               → Best for radial operations (derivatives, boundary conditions)

Spectral Space Pencil ((lmax+1) × (mmax+1) × nr)

The spectral pencil stores spherical harmonic coefficients on a rectangular degree/order grid. Invalid (l,m) slots such as m > l are left unmapped.

Arguments

• nlat, nlon, nr: Grid dimensions
• sht_config: SHTnsKit configuration (provides nlm)
• comm: MPI communicator
• optimize: accepted for API compatibility but IGNORED in Phase 1 (the 1D-θ grid is fixed (nprocs, 1)); Phase 2 (r×θ) will honor it via optimize_process_topology_shtnskit.

Returns

NamedTuple with pencil configurations: (:theta, :θ, :phi, :φ, :r, :spec, :mixed, :thetaphys). `thetaphysis a 2D(nlat, nlon)θ-distributed / φ-local prototype pencil whose θ-split matchespencils.r`, useful for structural invariant checks.

create_pencil_fft_plans(pencils, dims) -> Dict{Symbol,Any}

Create precomputed FFTW plans for efficient FFT operations.

Why Precomputed Plans?

FFTW achieves best performance when FFT plans are created once and reused. Plan creation involves:

1. Analyzing the input size and memory layout
2. Choosing optimal algorithm (Cooley-Tukey, Bluestein, etc.)
3. Possibly benchmarking different strategies

Plan Types Created

• :phi_forward / :phi_backward: FFT/IFFT along longitude (dimension 2)
• :theta_forward / :theta_backward: FFT/IFFT for theta pencil orientation

Technical Notes

• Plans operate on the parent() array of PencilArrays (the underlying Julia array)
• We use FFTW directly rather than PencilFFTPlan because we need single-dimension transforms, not full multi-dimensional distributed FFTs
• The plans are tied to specific array sizes and memory layouts

Arguments

• pencils: NamedTuple of Pencil configurations
• dims: Tuple (nlat, nlon, nr) of grid dimensions

Returns

Dict mapping plan names to FFTW plan objects. Contains :fallback => true on error.

create_scalar_matrices(config, domain, diffusivity, dt; scalar_bc_code, theta, T,
                       inner_regularity)

Create implicit scalar time-stepping matrices with Dirichlet/Neumann boundary conditions embedded in the first and last radial rows.

When inner_regularity = true the inner boundary row implements the ball centre regularity condition Θ′(r₁) = l·Θ(r₁)/r₁ (Robin with β = l; reduces to Neumann Θ′=0 for l=0). This is appropriate for full-sphere geometry where there is no physical inner boundary; the outer row is still determined by scalar_bc_code. Defaults to false (shell geometry, standard BCs on both ends).

create_shtns_aware_output_config(shtns_config, pencils; base_config=default_config())

Return an OutputConfig whose spectral output limit follows the supplied SHTns configuration.

The pencils argument is currently accepted to mirror setup call sites that construct output configuration together with decomposition metadata.

create_shtns_spectral_field(T, config, outer_core_domain, pencil_spec) -> SHTnsSpecField{T}

Create a new spectral field initialized to zero with default Dirichlet BCs.

Arguments

• T: Element type (typically Float64)
• config: SHTnsKit configuration providing nlm and other parameters
• outer_core_domain: RadialDomain specifying the radial discretization
• pencil_spec: PencilArrays Pencil defining the data distribution

Returns

A new SHTnsSpecField with:

• All spectral coefficients initialized to zero
• Dirichlet boundary conditions on all modes
• Zero boundary values

create_shtnskit_transpose_plans(pencils) -> Dict{Symbol,Any}

Create transpose plans for redistributing data between pencil orientations.

What is a Transpose?

In pencil decomposition, a "transpose" is an MPI all-to-all communication that redistributes data so a different dimension becomes local. For example:

• theta → phi transpose: makes longitude local (for FFTs)
• phi → r transpose: makes radius local (for radial derivatives)

Why Plans?

Like FFT plans, transpose plans encode the communication pattern and can optimize buffer allocation and MPI operations.

Transpose Constraints

PencilArrays requires that source and destination pencils differ in at most one distributed dimension. If they differ in two dimensions, we need a multi-step transpose (e.g., phi → theta → r).

Plan Keys Created

• :theta_to_phi, :phi_to_theta: Switch between theta and phi local
• :theta_to_r, :r_to_theta: Switch between theta and r local
• :phi_to_r, :r_to_phi: May be direct or marked as :multi_step

Usage

# Transpose data from theta-local to phi-local orientation
mul!(dest_array, transpose_plans[:theta_to_phi], src_array)

create_solver_erk2_cache(T, config, domain, diffusivity, dt; bc_spec=nothing, ...)

Precompute generic ERK2 propagators for velocity-like spectral fields.

create_solver_erk2_magnetic_poloidal_cache(T, config, domain, diffusivity, dt; ...)

Precompute ERK2 propagators for magnetic poloidal fields with embedded insulating boundary rows.

create_solver_erk2_magnetic_toroidal_cache(T, config, domain, diffusivity, dt; ...)

Precompute ERK2 propagators for magnetic toroidal fields with embedded homogeneous Dirichlet boundary rows.

create_solver_erk2_scalar_cache(T, config, domain, diffusivity, dt, boundary_condition; ...)

Precompute ERK2 propagators for scalar fields with embedded boundary rows.

The cache stores one set of matrices per unique spherical-harmonic degree in config. Dense matrices are precomputed unless use_krylov=true, in which case operator matrices are stored for Krylov actions.

create_solver_fields(T, backend)

Allocate the spectral/physical field containers required by the solver runtime for element type T.

create_solver_implicit_matrices(T, backend)

Precompute the linear solve operators used by the solver timestep schemes for element type T.

The returned matrices are later wrapped into the solver-owned matrix store so the timestep loop can reuse them without rebuilding operators on each step.

Returns (matrices, magnetic_ic_admittance). magnetic_ic_admittance is a NamedTuple{(:tor,:pol)} of InnerCoreAdmittance objects when the conducting inner-core path is enabled (params.magnetic_inner_bc === :conducting_inner_core) and nothing otherwise. The insulating default path is byte-for-byte unchanged.

create_solver_parameters([params])

Return a SolverParameters object for solver initialization.

create_solver_runtime(T, backend; auto_optimize=false, adaptive_threading=false)

Allocate the solver-owned field objects, shared workspaces, timestep state, and runtime domain references for a new run.

Scalar parameter BCs are installed before optional file BCs are loaded so a user-supplied spectral boundary file cleanly overrides the homogeneous or mean-mode defaults.

create_solver_velocity_poloidal_influence_matrices(T, config, domain, diffusivity, dt, velocity_bc_code; theta)

Build Green-function influence operators used to correct velocity-poloidal endpoint constraints after ERK2 finalization.

GeoDynamo.create_stress_free_pol_bc(T, d2_row)

Create a public poloidal-velocity stress-free endpoint descriptor.

GeoDynamo.create_stress_free_tor_bc(T, d1_row, r_inv)

Create a public toroidal-velocity stress-free endpoint descriptor.

create_temperature_matrices(config, domain, diffusivity, dt;
                             temperature_bc_code, theta, T, inner_regularity)

Create implicit time-stepping matrices for the temperature equation with Dirichlet/Neumann boundary rows embedded in the radial system matrix.

Pass inner_regularity = true for ball (full-sphere) geometry to impose the centre regularity condition Θ′(r₁) = l·Θ(r₁)/r₁ on the inner boundary row.

GeoDynamo.create_velocity_poloidal_influence_matrices(T, config, domain, diffusivity, dt, velocity_bcs; theta=0.5)

Public wrapper for building velocity-poloidal influence correction operators.

create_velocity_poloidal_matrices(config, domain, diffusivity, dt;
                                  theta, velocity_bc_code, T)

Create implicit time-stepping matrices for the poloidal velocity component with boundary conditions embedded in the matrix rows.

The boundary rows of the system matrix are replaced with the BC equations:

• No-slip: first derivative row (∂P/∂r = value)
• Stress-free: second derivative row (∂²P/∂r² = value)

create_velocity_poloidal_split_matrices(config, domain, Ek, dt; velocity_bc_code, theta=0.5, ball=false, T=Float64)

Stage-4B W-split operators for the pressure-free poloidal momentum equation Ek(∂t − Dpol)W = NW, W = Dpol·P, Dpol = ∂rr − l(l+1)/r² with P-recovery using Dirichlet rows at both walls (shell) or a Robin center-regularity inner row (ball=true), and cached no-slip influence responses enforcing P′=0 (velocitybc_code == 1). Other BC codes error loudly until ported. See docs/superpowers/plans/2026-06-10-double-curl-stage4b-*.md.

Ball mode (ball = true, full sphere with the off-center grid r₁ > 0): the inner wall conditions become center-regularity Robin rows exact for the leading behaviors P ~ r^{l+1} and W ~ r^{l+1}. The P-recovery inner row is the Robin P′(r₁) = (l+1)P(r₁)/r₁ (outer Dirichlet wall unchanged), and the influence matrix mixes spaces: row 1 = the W-regularity residual W′(r₁) − (l+1)W(r₁)/r₁ applied to the W-space Green columns gi (the condition lives on W = Dpol·P), row 2 = the outer wall residual on the recovered P responses hi (as shell). `d1rowinnerthen holds the PURE first-derivative endpoint row; the l-dependent −(l+1)/r₁ correction is applied at dot time viaregr_inv`.

create_velocity_toroidal_matrices(config, domain, diffusivity, dt;
                                  theta, velocity_bc_code, mass_coeff, T,
                                  inner_regularity)

Create implicit time-stepping matrices for the toroidal velocity component with boundary conditions embedded in the matrix rows.

The boundary rows of the system matrix are replaced with the BC equations:

• No-slip: identity row (T = value)
• Stress-free: ∂T/∂r - T/r = 0

When inner_regularity = true (ball / full-sphere geometry) the inner row instead imposes the center regularity condition t′(r₁) = l·t(r₁)/r₁ (β = l, raw sphtor toroidal scalar t ~ r^l); the outer row is still set by velocity_bc_code.

This ensures the implicit solve enforces BCs exactly rather than applying them as post-processing.

dense_to_cpu_spectral!(field_spec, dense_r, dense_i, config, nr) -> field_spec

Inverse of cpu_spectral_to_dense: scatter dense (lmax+1, mmax+1, nr) real/imag arrays back into the CPU slot-packed spectral storage of field_spec ((l+1, m+1) → mode slot). dense_* may be host or device arrays (host-copied here).

GeoDynamo.disable_erk2_diagnostics!()

Disable ERK2 residual diagnostics without changing the configured interval.

GeoDynamo.enable_erk2_diagnostics!(; interval=...)

Enable ERK2 residual diagnostics and optionally update the reporting interval.

GeoDynamo.enforce_erk2_bc!(result, bc_side, boundary_idx, l, nr; value_override=nothing)

Public wrapper for enforcing one ERK2 boundary condition on a radial profile.

enforce_velocity_boundary_values!(𝒰)

Anchor toroidal and poloidal spectral coefficients to the currently cached Dirichlet boundary values on the inner and outer radial surfaces.

GeoDynamo.erk2_apply_stage!(buffers, u)

Public wrapper that writes the provisional ERK2 stage into a field.

GeoDynamo.erk2_diagnostics_enabled()

Return whether ERK2 residual diagnostics are currently enabled.

GeoDynamo.erk2_diagnostics_interval()

Return the configured ERK2 residual diagnostics interval.

GeoDynamo.erk2_finalize_field!(buffers, u, cache, config, dt; bc_spec=nothing)

Public wrapper for writing the accepted ERK2 update back into a field.

GeoDynamo.erk2_prepare_field!(buffers, u, nl, cache, config, dt; bc_spec=nothing)

Public wrapper for preparing the provisional ERK2 stage for one field.

GeoDynamo.erk2_stage_residual_stats(buffers)

Public wrapper returning global max and L2 ERK2 stage-residual norms.

GeoDynamo.erk2_store_stage_nonlinear!(buffers, nl)

Public wrapper for storing nonlinear terms evaluated at the ERK2 stage.

estimate_field_count() -> Int

Estimate the number of field arrays typically allocated simultaneously. Used for memory usage estimation. Returns 6 as a reasonable default covering velocity (3 components), temperature, composition, and magnetic field.

estimate_memory_usage(pencils, field_count::Int, precision::Type)

Estimate memory usage for given pencil configuration.

estimate_memory_usage_shtnskit(nlat, nlon, nr, lmax, field_count, T)

Estimate memory usage for SHTnsKit-based transforms with PencilArrays.

Arguments

• nlat: Number of latitude points
• nlon: Number of longitude points
• nr: Number of radial points
• lmax: Maximum spherical harmonic degree
• field_count: Number of fields to estimate for
• T: Element type (e.g., Float64)

evaluate_chebyshev_derivative(n::Int, r::T, domain::RadialDomain) where T

Evaluate derivative of Chebyshev polynomial at a point.

evaluate_tau_derivative_inner(tau::Vector, domain::RadialDomain)

Evaluate tau function derivative at inner boundary.

evaluate_tau_derivative_outer(tau::Vector, domain::RadialDomain)

Evaluate tau function derivative at outer boundary.

exp_action_krylov(Aop!, v, dt; m=20, tol=1e-8)

Compute exp(dt*A)v with an Arnoldi/Krylov approximation.

The high-level timestep drivers call exp_action_krylov directly, but this root-level entry point remains available as a reusable numerical utility.

extract_all_fields(state::SolverState)

Return a copy-based dictionary snapshot of the main solver fields.

This is primarily used by restart/output tooling and tests that need a stable container representation independent of the in-memory field types.

extract_coefficients_for_shtnskit!(coeffs_buffer, spec_real, spec_imag, r_local, config)

Extract spectral coefficients into SHTnsKit's expected matrix format.

Data Format Conversion

Our internal format: Linear array indexed by combined (l,m) index SHTnsKit format: Matrix indexed by [l+1, m+1]

Arguments

• coeffs_buffer: Pre-allocated (lmax+1) × (mmax+1) complex matrix (output)
• spec_real, spec_imag: Real and imaginary parts of spectral coefficients
• r_local: Radial level index (1-based)
• config: SHTnsKit configuration

Threading

Uses @threads for parallel filling of the coefficient matrix.

extract_coefficients_for_shtnskit(spec_real, spec_imag, r_local, config) -> Matrix{ComplexF64}

High-level coefficient extraction with automatic buffer management and MPI gathering.

This is the main entry point for preparing spectral coefficients for SHTnsKit. It handles:

1. Buffer allocation/reuse from cache
2. Local coefficient extraction
3. MPI Allreduce to combine coefficients from all processes

Why MPI Gathering?

Spectral coefficients may be distributed across MPI processes. SHTnsKit needs the complete coefficient matrix, so we sum partial contributions from all processes.

Returns

A complete (lmax+1) × (mmax+1) coefficient matrix ready for SHTnsKit.synthesis()

extract_coefficients_pair_for_shtnskit(spec1_real, spec1_imag, spec2_real, spec2_imag, r_local, config)

Extract two spectral coefficient matrices efficiently for vector transforms.

This is optimized for the common case in vector synthesis/analysis where we need both toroidal and poloidal coefficients. It avoids one copy operation compared to calling extract_coefficients_for_shtnskit twice.

Returns

Tuple (coeffs1, coeffs2) of coefficient matrices.

extract_field_info(fields[, config, pencils])

Infer a FieldInfo description from a dictionary-style field snapshot.

This bridges the simulation-side field containers and the output writer’s more uniform metadata needs.

extract_local_radial_profile!(profile, data, slot, nr, r_range)

In-place version to avoid allocations; writes the local radial line into profile for the given local spectral slot using the provided r_range.

extract_physical_slice_generic!(slice_buffer, phys_data, r_local, config)

Generic extraction for any pencil orientation using pre-allocated buffer.

WARNING: MPI Synchronization

This function contains MPI.Allreduce! which is a collective operation. When called inside a per-radial loop, ALL MPI processes must call this function the same number of times, otherwise deadlock will occur. Ensure even radial distribution or use global loop bounds.

extract_physical_slice_phi_local!(slice_buffer, phys_data, r_local, config)

Extract physical slice when in phi-local pencil using pre-allocated buffer.

WARNING: MPI Synchronization

This function contains MPI.Allreduce! which is a collective operation. When called inside a per-radial loop, ALL MPI processes must call this function the same number of times, otherwise deadlock will occur. Ensure even radial distribution or use global loop bounds.

field_to_device(arch, host_phys::AbstractArray, config, nr) -> GPUPhysicalField
field_to_device(arch, (hr, hi)::Tuple, config, nr)          -> GPUSpectralField

Copy host data onto arch's backend, wrapped in the matching GPU field.

field_to_host(f) -> NamedTuple

Copy a GPU field's device arrays back to host Arrays. Returns (; data) for a GPUPhysicalField, (; data_real, data_imag) for a GPUSpectralField.

finalize_solver_erk2_field!(buffers, u, cache, config, dt; bc_spec=nothing)

Write the accepted ERK2 update back into u.

The final state combines the full-step linear propagation, the first nonlinear increment, and the phi2 correction from the staged nonlinear residual.

finalize_solver_step!(state, step)

Advance the runtime clock after a timestep and refresh solver views.

find_package_root()

Find the root directory of the GeoDynamo.jl package.

find_restart_files(restart_dir, target_time)

Return restart NetCDF files ordered best-match first.

target_time <= 0 is the "restart from the latest checkpoint" sentinel and orders files newest-first by modification time (cheap; no files opened). A positive target_time is honored: each checkpoint's stored simulation time is read and files are ordered by closeness to target_time, so a requested restart time actually selects the nearest snapshot rather than the newest file. Files whose stored time cannot be read fall back to the end (newest-first).

force_curl_projections!(R_tor, R_pol, Q, S, T_in, domain)

Momentum-equation projections of a force field from its QST scalars (see force_physical_to_qst!):

R_tor[lm](r) = [r̂·∇×F]_lm          = −λ/r · T_lm(r)
R_pol[lm](r) = [r̂·∇×∇×F]_lm = (λ/r²) · (Q_lm(r) − ∂_r(r · S_lm(r)))

where λ = l(l+1). Verified against closed-form analytic references (pure-radial, spheroidal, toroidal families) and an exact single-curl numerical reference in test/forceprojectionreference.jl.

force_physical_to_qst!(force, Q, S, T)

Three-component spectral analysis of a (generally non-solenoidal) force field. Q <- scalar analysis of the radial component; (S, T) <- spheroidal/toroidal scalars of the tangential components (the same sphtor analysis the velocity path uses – S is what that path stores in its "poloidal" output).

from_sht_spectral(cfg, alm::AbstractMatrix{ComplexF64})

dist_analysis already returns the complete dense matrix replicated on every rank, so no gather is needed — return it (identity). Symmetric with to_sht_spectral.

from_spec_solve!(cfg, Alm, solve, plan)

Transpose the solver-orientation PencilArray solve (output of the radial implicit step) back into Alm (the DistTransposePlan layout). Overwrites Alm in-place.

solve MUST be scratch.solve (the adapter's cached solve buffer, returned by to_spec_solve or spec_storage_to_solve!). The persistent Transposition plan scratch.t_bwd is bound to that specific array, allowing zero-alloc repeated use.

generate_filename(config, time, step, file_type="output", output_number=1; geometry=:shell)

Build the on-disk NetCDF path for history, restart, or auxiliary output files.

Files are numbered by writer counters rather than by time or step, which keeps names stable across floating-point time representations.

get_backend(arch)

Return the backend object associated with arch.

get_bc_vectors(field)

Return mode-indexed scalar boundary vectors for timestep solves.

When a spectral BC file has been loaded, the interpolation cache supplies real and imaginary values for each mode. Otherwise the solver falls back to the field's parameter-derived boundary_values, which contains real values only.

Always returns a _BCVectors with the four keys inner_real, outer_real, inner_imag, outer_imag; absent slots are nothing.

get_default_nlat() -> Int

Return the fallback number of latitude points (64) used when no grid is available at call time (e.g., during precompilation).

get_default_nlon() -> Int

Return the fallback number of longitude points (128) used when no grid is available at call time (e.g., during precompilation). Power of 2 for FFT.

get_disttranspose_plan(cfg) -> SHTnsKit.DistTransposePlan

Return the DistTransposePlan associated with cfg, building and caching it on the first call. Subsequent calls are O(1) field reads.

The plan is stored in cfg._buffers.disttranspose_plan (a mutable holder). The plan is constructed on cfg.pencils.θ_comm with nlev = nr_local.

get_file_info(filename)

Return basic metadata for one NetCDF output file: time, step, dimensions, and variable names.

get_flux_value(lm_idx::Int, boundary::Int, 𝔽::AbstractScalarField)

Get prescribed flux value for given mode and boundary from scalar field. This is a generic version that works with any scalar field.

get_pencil_orientation(pencil) -> Symbol

Determine which dimension(s) are fully local in a pencil decomposition.

Returns

• :theta_phi: Both angular dimensions local (serial or single-node)
• :theta: Latitude (θ) dimension is local
• :phi: Longitude (φ) dimension is local
• :r: Radial dimension is local

Usage

Used to choose optimal transform strategies based on data layout.

get_solver_erk2_cache!(caches, key, diffusivity, T, config, domain, dt; ...)

Compatibility shim that dispatches legacy Symbol keys to the concrete velocity cache overloads.

get_solver_erk2_cache!(caches, Val(:velocity_poloidal), diffusivity, T, config, domain, dt; ...)

Return the velocity-poloidal ERK2 cache from TimestepCaches.

get_solver_erk2_cache!(caches, Val(:velocity_toroidal), diffusivity, T, config, domain, dt; ...)

Return the velocity-toroidal ERK2 cache from TimestepCaches.

This concrete overload avoids runtime Symbol dispatch in the main solver step while still sharing the generic rebuild checks.

get_solver_erk2_composition_cache!(caches, diffusivity, T, config, domain, dt, composition_bc_code; ...)

Return the solver-owned composition ERK2 cache with the same compatibility checks used for the temperature scalar cache.

get_solver_erk2_influence_matrices!(cache, key, T, config, domain, diffusivity, dt, velocity_bc_code; theta)

Return velocity-poloidal influence matrices from TimestepCaches.

Only :velocity_poloidal is valid here because the correction enforces the poloidal boundary constraints after the ERK2 field update.

get_solver_erk2_magnetic_poloidal_cache!(caches, diffusivity, T, config, domain, dt; ...)

Return or rebuild the magnetic-poloidal ERK2 cache stored in TimestepCaches.

get_solver_erk2_magnetic_toroidal_cache!(caches, diffusivity, T, config, domain, dt; ...)

Return or rebuild the magnetic-toroidal ERK2 cache stored in TimestepCaches.

get_solver_erk2_temperature_cache!(caches, diffusivity, T, config, domain, dt, temperature_bc_code; ...)

Return the solver-owned temperature ERK2 cache, rebuilding it when timestep, grid, diffusivity, Krylov settings, or boundary conditions no longer match.

get_time_series(output_dir)

Read and return sorted simulation times from history NetCDF files in output_dir.

Unreadable files are skipped so a partially written or unrelated file does not break post-processing.

gpu_apply_bc_rows!(x_r, x_i, bc_in_r, bc_in_i, bc_out_r, bc_out_i) -> nothing

Overwrite the boundary rows of the per-mode radial RHS with the prescribed BC values: row 1 (inner) ← bc_in_*, row nr (outer) ← bc_out_*. x_* are (nl,nm,nr); bc_* are (nl,nm) (per-(l,m) boundary value).

gpu_batched_banded_matvec!(Y, X, mat, bw) -> Y

Apply the banded radial operator mat ((2bw+1,nr)) to every mode's radial profile: Y[l,m,:] = mat · X[l,m,:]. Y/X are (nl,nm,nr). Y must NOT alias X (an output point reads input points at other radii). Backend inferred from Y.

gpu_batched_banded_matvec_perl!(Y, X, mat_batched, bw) -> Y

Per-degree-l banded mat-vec: Y[l,m,:] = mat_batched[:,:,l] · X[l,m,:], where mat_batched is (2bw+1, nr, nl) (degree l = dim-3 slice). Y/X are (nl,nm,nr); Y must NOT alias X. Backend inferred from Y.

gpu_batched_banded_solve!(X, B, lu_batched, bw) -> X

Solve A_l · X[l,m,:] = B[l,m,:] for every (l,m), where A_l's banded LU is lu_batched[:,:,l] (degree l = dim-3 index; lu_batched is (2bw+1,nr,nl)). X/B are (nl,nm,nr); in-place X === B is supported. Backend (CPU/CUDA) is inferred from X.

gpu_batched_dense_matvec_perl!(Y, X, M3) -> Y

Apply the per-degree DENSE radial operator stack M3 :: (nr, nr, nl) to every mode's radial profile: Y[l,m,:] = M3[:,:,l+1] · X[l,m,:]. Y must not alias X.

gpu_build_rhs_cnab2!(rr, ri, ur, ui, nr_, ni_, pr, pi_, lin_batched, inv_dt, linear_weight, bw) -> nothing

Assemble the CNAB2 implicit RHS (split real/imag): rhs = inv_dt·u + 1.5·nl − 0.5·nl_prev + linear_weight·(lin·u), where lin·u is the per-l banded mat-vec of the linear operator (lin_batched (2bw+1,nr,nl)). inv_dt = mass_coeff/dt, linear_weight = 1−θ. All arrays (nl,nm,nr) on the same backend; outputs distinct from inputs. (pi_ is the previous-step imaginary nonlinear term — the trailing underscore avoids shadowing Julia's pi constant.)

gpu_buoyancy_add!(force_r, s, r_vec, factor) -> force_r

Add the radial buoyancy/codensity force force_r += factor·r·s, with r per radial level (r_vec length-nr, radial = dim 3). Use factor=(Pm/Pr)·Ra for thermal buoyancy or factor=(Pm/Sc)·Ra_C for compositional (matching the CPU). All arrays (including r_vec) must live on the same backend — mixing host and device arrays errors at broadcast time.

gpu_cb3_solver_step!(state) -> nothing

Advance the dense GPU solver state by one RungeKutta3 IMEX-RK3 step. state must come from build_gpu_solver_state(cpu_state) where cpu_state.parameters.timestepper isa RungeKutta3.

gpu_coriolis_sub!(out_r, out_θ, out_φ, u_r, u_θ, u_φ, sinθ, cosθ) -> nothing

Subtract the Coriolis term ẑ×u from the accumulator (CPU: adv_i -= (ẑ×u)_i), where (ẑ×u)_r=-sinθ·u_φ, (ẑ×u)_θ=-cosθ·u_φ, (ẑ×u)_φ=cosθ·u_θ+sinθ·u_r. The net effect is out_r += sinθ·u_φ, out_θ += cosθ·u_φ, out_φ -= cosθ·u_θ+sinθ·u_r. sinθ,cosθ are length-nlat (latitude = dim 1). The 2Ω factor is absorbed into the nondimensional coefficients upstream (Ekman number), matching the CPU. All arrays (including sinθ/cosθ) must live on the same backend — mixing host and device arrays errors at broadcast time.

gpu_cross!(out_r, out_θ, out_φ, a_r, a_θ, a_φ, b_r, b_θ, b_φ, coeff) -> nothing

Overwrite out = coeff·(a×b) component-wise. Cross-product order matches the CPU Lorentz (coeff=1/Pm), velocity advection (coeff=E), and induction (coeff=1).

gpu_cross_add!(out_r, out_θ, out_φ, a_r, a_θ, a_φ, b_r, b_θ, b_φ, coeff) -> nothing

Accumulate out += coeff·(a×b) component-wise (e.g. add the Lorentz force onto an advection accumulator).

gpu_erk2_solver_step!(state, erk) -> nothing

Advance every field one staged exponential RK2 step on the dense device state state (from build_gpu_solver_state) with the ERK2 operator pack erk (from build_gpu_erk2_state). Mirrors the CPU compute_solver_nonlinear_terms! + integrate_solver_erk2_step! sequence, including the velocity-poloidal W-split V-advance with φ1-column P-recovery at the stage and the finalize. The lagged physical buffers evolve exactly as on the CPU (each nonlinear pass refreshes them).

gpu_functional() -> Bool

true only when a CUDA-capable GPU is present AND the GeoDynamoCUDAExt extension is loaded (i.e. CUDA.functional()). false otherwise, including on machines with no GPU. Use this to gate GPU code paths and tests.

gpu_implicit_solve_field!(x_r, x_i, lu_batched, bc_in_r, bc_in_i, bc_out_r, bc_out_i, bw) -> nothing

One field's CNAB2 implicit solve: x holds the RHS on entry; this sets the BC boundary rows (gpu_apply_bc_rows!) then batched-solves the per-mode banded systems in place (Phase 4 gpu_batched_banded_solve!, X===B supported), leaving the solution in x. lu_batched (2bw+1,nr,nl) are the per-l LU factors.

gpu_inner_core_history_flux!(φ0_r, φ0_i, S_old_r, S_old_i, ic) -> nothing

Per-mode conducting-inner-core CNAB2 history flux φ0 = d1_top·y, where M_ic y = b with b = inv_dt·S_old + weight·L·S_old and homogeneous boundary rows (b[1]=b[Nic]=0). S_old_* are dense (nl,nm,Nic) inner-core spectra; φ0_* are (nl,nm). ic is the bundle from gpu_pack_inner_core. Mirrors inner_core_history_flux (inner_core.jl:168-175). All arrays on the same backend.

gpu_magnetic_field_step!(tor, pol, u_r, u_θ, u_φ, config, nlops, inv_dt, linear_weight,
                         lmax, bw; continuity_mag=false, ic=nothing) -> nothing

Advance the magnetic field one CNAB2 step. tor/pol are NamedTuple bundles (; spec_r, spec_i, prev_nl_r, prev_nl_i, lin, lu); on exit *.spec_* is the updated field and *.prev_nl_* holds THIS step's nonlinear term. u_* is the physical velocity (supplied — from the velocity step). nlops carries the magnetic nonlinear/curl operators. inv_dt = 1/dt and lin carry the magnetic mass coefficient (η is in lin); linear_weight = 1−θ.

Insulating inner core (default, ic = nothing): continuity_mag=true applies the CONTINUITY_MAG toroidal inner-boundary coupling: the toroidal inner RHS row is set to −nl_pol[ICB] + prev_nl_pol[ICB] (ICB = radial index 1), computed from the just-formed poloidal nonlinear and the OLD poloidal history. Otherwise the toroidal inner row is 0. The poloidal is fully homogeneous.

Conducting inner core (ic keyword given): ic must be a NamedTuple (; tor_adm, pol_adm, tor_ic_r, tor_ic_i, pol_ic_r, pol_ic_i) where *_adm are gpu_pack_inner_core bundles and *_ic_* are (nl, nm, Nic) inner-core spectral state arrays. When ic is supplied:

• The inner BC for BOTH toroidal and poloidal is the conducting history flux φ0 from gpu_inner_core_history_flux! (5l), computed from the OLD inner-core state BEFORE the outer-core solve. This is NOT incremental (no prev subtraction, unlike the CONTINUITYMAG path); `continuitymag` is ignored.
• Outer BC = 0 (same as insulating).
• After the outer-core field update, the inner-core profile is reconstructed via gpu_reconstruct_inner_core! (5l) using g = spec[:,:,1] (the NEW outer-core ICB value, materialized — not a view), into fresh scratch, then written back into the ic.*_ic_* arrays in place.

ORDERING INVARIANT: φ0 reads OLD ic state; both build_rhs calls and the CONTINUITYMAG BC all read OLD outer-core state (`*.spec,pol.prevnl`); the field/history and ic state are overwritten ONLY after every such read. All arrays on the same backend.

gpu_magnetic_nonlinear!(nl_tor_r, nl_tor_i, nl_pol_r, nl_pol_i, B_tor_r, B_tor_i, B_pol_r, B_pol_i,
                        u_r, u_θ, u_φ, config, d1, d2, lfac, rinv, rinv2, rscale, lmax, bw;
                        r_vec=nothing) -> nothing

Magnetic induction nonlinear nl = ∇×(u×B). B_tor/B_pol the magnetic toroidal/poloidal spectral; u_* the physical velocity (supplied); nl_tor/nl_pol the toroidal/poloidal induction nonlinear. d1/d2/lfac/rinv/rinv2 are the curl/transform operators; rscale=1/r² is the radial-synthesis scale forwarded to gpu_vector_spectral_to_physical! (Stage-2 B_r = l(l+1)·P·rscale) — LOAD-BEARING, not optional. All on the same backend; outputs distinct from inputs.

gpu_pack_banded_lu(lus, arch) -> (2bw+1, nr, nl) array on arch's backend

Stack the per-degree BandedLU factors (lus[l] for l-slot l) into a single batched array lu_batched[:, :, l] = lus[l].lu, on arch's backend. All factors must share the same bandwidth bw and size nr.

gpu_pack_influence(influence, nl, nr, arch) -> (Gre_b, invG_b)

Flatten the per-degree ERK2InfluenceOp correction operators into batched arrays on arch's backend: Gre_b is (nr,2,nl) and invG_b is (2,2,nl), indexed by dim-3 = dense degree slot li (degree l = li-1). Degrees absent from influence (including l=0) get all-zero columns, so the kernel applies an exact no-op to those modes — matching the CPU path, which skips them.

influence is the Dict{Int,ERK2InfluenceOp{T}} keyed by degree l (0-based). nl is the number of degree slots (lmax+1); nr the radial size.

gpu_pack_inner_core(adm::InnerCoreAdmittance, nl, arch)
    -> (; lin_ic, lu_ic, d1_top, inv_dt, weight, Nic, bw)

Flatten the per-degree conducting-inner-core operators into batched arrays on arch's backend: lin_ic / lu_ic are (2bw+1, Nic, nl) indexed by dim-3 = degree slot li (degree l = li-1). Stored degrees carry adm.lin[…].data and adm.factor[…].lu; non-stored degrees (including l=0) get a zero L and an identity LU (diagonal row = 1) so the batched solve treats those empty modes as x = b rather than dividing by a zero pivot. Returns the operators plus the CNAB2 scalars inv_dt = 1/dt, weight = 1−θ, and Nic/bw.

gpu_phi_gradient!(gφ_r, gφ_i, s_r, s_i, mvals) -> nothing

Longitudinal gradient ∂s/∂φ = i·m·s: gφ_r = −m·s_i, gφ_i = m·s_r. mvals[m+1] = m (length nm, m-slot index).

gpu_poloidal_wsplit_advance!(P_new, P, nl, prev_nl, ws, inv_dt, linear_weight, bw) -> P_new

One CNAB2 step of the Stage-4B poloidal W-split for ONE split-complex component (call once for real, once for imag), batched over all (l,m) modes. Mirrors the CPU _apply_poloidal_wsplit_cnab2! (physics/velocity/solver.jl):

1. W = D_pol·P, LW = Ek·D_pol·W
2. rhs = (Ek/dt)·W + (1−θ)·LW + 1.5·nl − 0.5·prev_nl (nl carries N_W)
3. W⁺ = A_W⁻¹·rhs (banded LU per degree)
4. Dirichlet P-recovery: zero W⁺ endpoint rows, P⁺ = A_P⁻¹·W⁺
5. endpoint-residual influence correction: ρᵢ = d1_rowᵢ·P⁺, solve the per-degree 2×2 M·a = −ρ, P_new = P⁺ + a₁·h₁ + a₂·h₂; the l = 0 plane carries no poloidal flow.

ws is the device-state wsplit operator pack (see _build_wsplit_pack). inv_dt = Ek/dt (the same mass-scaled value the toroidal CNAB2 uses) and linear_weight = 1−θ. P_new must not alias P/nl/prev_nl.

gpu_reconstruct_inner_core!(S_new_r, S_new_i, S_old_r, S_old_i, g_r, g_i, ic) -> nothing

Per-mode conducting-inner-core reconstruction: solve M_ic S = b with b = inv_dt·S_old + weight·L·S_old, regularity b[1]=0, and ICB Dirichlet b[Nic]=g (the outer-core value at the ICB). S_old_* dense (nl,nm,Nic), g_* (nl,nm); the solution is written to S_new_*. Mirrors reconstruct_inner_core (innercore.jl:185-192). `Snew*may not aliasSold_*`. All arrays on the same backend.

gpu_run!(state, nsteps; output_every = 0, output_fn = nothing) -> state

Advance the device solver state (from build_gpu_solver_state, optionally moved to a device via gpu_to_device) by nsteps CNAB2 steps, in place.

Each iteration calls gpu_solver_step!; the per-field prev_nl rollover and the persistent lagged physical buffers are mutated in place inside the step, so the across-step history evolves correctly without extra bookkeeping here.

If output_fn !== nothing and output_every > 0, then after every output_every-th step output_fn(host_state, step) is called, where host_state = gpu_to_device(state, CPU()) is a host-gathered deep copy (device → host dense Arrays) suitable for IO / restart writing without disturbing the live (possibly device-resident) state.

gpu_run!(cpu_state::SolverState, nsteps; arch = CPU(),
         output_every = 0, output_fn = nothing) -> cpu_state

Public-API convenience: advance a configured CPU SolverState by nsteps CNAB2 steps ON THE GPU PATH. Builds the device state via build_gpu_solver_state (optionally moved to arch with gpu_to_device — pass arch = GPU() on a CUDA box), runs the device loop, then ALWAYS syncs the evolved state back into cpu_state via sync_gpu_state_to_cpu! (spectral fields, CNAB2 prev_nl histories, and the lagged physical buffers) and advances cpu_state.step/.time, so CPU-side stepping / diagnostics / output / restart can continue coherently from the GPU-evolved state. (To run the device loop without syncing back — keeping a handle to the device state — call build_gpu_solver_state + gpu_run!(gst, …) directly instead.)

cpu_state should be initialized (≥1 prior solver_step!, or initialize_solver_fields!) so the CNAB2 history (prev_nl) and the lagged physical buffers are populated, exactly as the CPU path requires. Builder scope is insulating magnetic + CNAB2 (see build_gpu_solver_state).

gpu_scalar_advection!(out, u_r, u_θ, u_φ, ∇r, ∇θ, ∇φ) -> out

out = -(u_r·∇r + u_θ·∇θ + u_φ·∇φ) — the scalar advection -(u·∇)s.

gpu_scalar_field_step!(spec_r, spec_i, prev_nl_r, prev_nl_i, u_r, u_θ, u_φ, config,
                       d1, mvals, rinv, lin_batched, lu_batched,
                       bc_in_r, bc_in_i, bc_out_r, bc_out_i, inv_dt, linear_weight, lmax, bw) -> nothing

Advance one scalar field one CNAB2 step. On entry spec_* is the field and prev_nl_* the previous nonlinear term; on exit spec_* is the updated field and prev_nl_* holds THIS step's nonlinear term (rolled over). lin_batched are the per-l linear operators L, lu_batched the per-l LU factors of the system matrix (I−θ·dt·L); bc_* the per-mode BC values; inv_dt = mass_coeff/dt, linear_weight = 1−θ. All arrays on the same backend.

gpu_scalar_gradient!(gr_r,gr_i, gθ_r,gθ_i, gφ_r,gφ_i, s_r,s_i,
                     d1, mvals, rinv, lmax, bw) -> nothing

Assemble the scalar gradient: ∇r = d1·s (banded mat-vec, NOT scaled by 1/r), ∇θ via the Legendre recurrence, ∇φ = i·m·s, then multiply the tangential components (∇θ,∇φ) by rinv = 1/r (0 at r=0). d1 banded (2bw+1,nr); mvals length-nm; rinv length-nr; all on the same backend; outputs distinct from s_r/s_i.

gpu_scalar_nonlinear!(nl_r, nl_i, s_r, s_i, u_r, u_θ, u_φ, config, d1, mvals, rinv, lmax, bw) -> nothing

Compute a scalar field's nonlinear term nl = analyze( −(u·∇s) ): gradient of s (spectral) → transform the gradient components to physical → advection −(u_r·∇r + u_θ·∇θ + u_φ·∇φ) against the supplied physical velocity → analyze the product back to spectral. nl_*/s_* are dense (nl,nm,nr); u_* physical (nlat,nlon,nr); d1/mvals/rinv as in gpu_scalar_gradient!. All on the same backend; nl_* distinct from s_*. (Scratch is pooled via GPUWorkspace when ws is supplied.)

gpu_scalar_physical_to_spectral!(spec, phys, config) -> spec

Analyze each radial level of the physical field phys ((nlat, nlon, nr)) into the dense spectral field spec ((lmax+1, mmax+1, nr) split real/imag), via SHTnsKit's per-level transform (GPU on CuArrays, CPU otherwise).

gpu_scalar_spectral_to_physical!(phys, spec, config) -> phys

Synthesize each radial level of the dense spectral field spec ((lmax+1, mmax+1, nr) split real/imag) into the physical field phys ((nlat, nlon, nr)), via SHTnsKit's per-level transform (GPU on CuArrays, CPU otherwise). config is the GeoDynamo SHTnsKitConfig.

gpu_scratch!(ws, key, ref, dims) -> array

Pooled similar(ref, dims) — same rules, explicit shape (e.g. one (nlat, nlon) level buffer derived from a 3-D field).

gpu_scratch!(ws, key, ref) -> array

Pooled similar(ref): first call under key allocates, later calls return the cached buffer (contents are stale — callers must overwrite). ws === nothing falls back to a plain similar(ref) (un-pooled paths keep working).

gpu_scratch_complex!(ws, key, ref_real, dims) -> dense complex buffer

Pooled complex buffer for the per-level transform staging (the complex.(view, view) materializations). ref_real supplies the backend.

gpu_scratch_phys!(ws, key, T, arch, config, nr) -> GPUPhysicalField

Pooled physical-field scratch (the (nlat, nlon, nr) work fields every transform/nonlinear pass needs). Same keying rules as gpu_scratch!. Returns are intentionally un-asserted (Array vs CuArray backing); call sites hand the field straight to kernel wrappers, which are dispatch barriers.

gpu_solver_step!(state) -> nothing

Advance every field one CNAB2 step on the GPU, in the CPU order velocity → magnetic → temperature → composition, with the shared physical velocity and the one-step physical-field lag for velocity's buoyancy/Lorentz coupling. See the file header for the state bundle layout and the lag semantics. state.magnetic / state.composition may be nothing to skip those fields. All arrays on the same backend.

gpu_spectral_curl!(dst_tor_r, dst_tor_i, dst_pol_r, dst_pol_i,
                   src_tor_r, src_tor_i, src_pol_r, src_pol_i,
                   d1, d2, lfac, rinv, rinv2, r_vec, bw) -> nothing

Spectral curl of a toroidal–poloidal field (vorticity ∇×u from velocity, or current ∇×B from magnetic — the same operator): dst_tor = rinv·(d2·P − lfac·rinv2·P), dst_pol = −r·T, with P=src_pol, T=src_tor. lfac[l+1]=l(l+1) (length nl); rinv/rinv2 length nr; r_vec is r; d1/d2 are banded (2bw+1,nr). d1 is retained in the signature for operator-bundle symmetry but is not used by the Stage-2 formula. All arrays on the same backend. Real/imag handled independently (curl is real-linear). The dst_* arrays must NOT alias any src_* array.

gpu_synchronize()

Block until all queued GPU work completes. No-op when no GPU backend is active.

gpu_theta_gradient!(gθ_r, gθ_i, s_r, s_i, lmax) -> nothing

Latitudinal gradient ∂s/∂θ via the Legendre (l±1, m) recurrence (matching the CPU). Empty slots (l < m) are zeroed. Outputs must be distinct from inputs (the recurrence reads neighbor l-slots). Backend inferred from gθ_r.

gpu_to_device(state, arch) -> NamedTuple

Deep-copy a build_gpu_solver_state bundle to arch's backend: every array (nested per-field bundles, nlops_*, influence, the physical buffers) is moved via on_architecture(arch, …); scalars, config, and nothing pass through. Use this to move a CPU-built state to GPU() for the GPU≈CPU hardware gate, then run gpu_solver_step! on the device copy. gpu_to_device(state, CPU()) returns an independent deep copy on the host.

gpu_vector_physical_to_spectral!(tor, pol, vθ, vφ, config;
                                 vr=nothing, lfac=nothing, r2=nothing,
                                 raw_spheroidal=false) -> nothing

Stage-2 analysis: per level analysis_sphtor(vθ, vφ) → raw (S, T); toroidal ← T. raw_spheroidal=true stores the raw S in the poloidal slot (the tangential-basis primitive). The default mode does the Q-based poloidal recovery (requires vr, lfac, r2): Q = scalar analysis of v_r; P = r²·Q/λ with λ = l(l+1) (l=0 slot zeroed) — the exact inverse of the solenoidal synthesis. Mirrors CPU vector_physical_to_spectral! + _poloidal_from_radial_q!.

gpu_vector_physical_to_spectral!(tor, pol, vr, vθ, vφ, config, lfac, rinv2) -> nothing

Analyze a solenoidal physical vector field under the Stage-2 convention. Toroidal coefficients come from raw tangential sphtor analysis; poloidal coefficients are recovered from radial scalar analysis, P = Q / (l(l+1)/r²).

gpu_vector_spectral_to_physical!(vr, vθ, vφ, tor, pol, config, lfac, rscale,
                                 d1, rinv, bw; raw_spheroidal=false) -> nothing

Stage-2 solenoidal synthesis: tangential (vθ, vφ) per level via synthesis_sphtor(S, T) with S = (d1·P)·(1/r); raw_spheroidal=true passes the stored poloidal coefficients directly as S (the tangential-basis primitive). Radial v_r per level via scalar synthesis of P·lfac·rscale with rscale = 1/r² (u_r = l(l+1)·P/r²). d1 is the banded ∂r operator ((2bw+1, nr)), rinv = 1/r. Mirrors the CPU vector_spectral_to_physical_disttranspose! convention. All operator arrays must be on the same backend as the fields (use on_architecture).

gpu_velocity_field_step!(tor, pol, config, nlops, influence, inv_dt, linear_weight,
                         lmax, bw; <coupling kwargs>) -> nothing

Advance the velocity field one CNAB2 step. tor/pol are NamedTuple bundles of the toroidal/poloidal arrays (see the file header); on entry *.spec_* is the field and *.prev_nl_* the previous nonlinear term, on exit *.spec_* is the updated field and *.prev_nl_* holds THIS step's nonlinear term (rolled over). nlops carries the nonlinear/curl operators, influence the poloidal 2×2 correction operators. inv_dt = E/dt and the per-l lin operators must already carry the mass coefficient E (so 5c reproduces the velocity RHS); linear_weight = 1−θ. Coupling kwargs (thermal/compositional buoyancy, Lorentz) are forwarded to gpu_velocity_nonlinear!; omit them for the velocity-only step. All arrays on the same backend.

ROTATING INNER CORE: the toroidal bc_in_*/bc_out_* rows are applied verbatim by the implicit solve; for a rotating inner core the caller must set the l=1, m=0 inner BC slot to the prescribed rotation value (rot_omega·r_inner, incremental) BEFORE calling — this function does not assemble it.

ORDERING INVARIANT (as in gpu_scalar_field_step!): the nonlinear and BOTH build_rhs calls read the OLD *.spec_*; the spec is overwritten with the solution ONLY after every such read. Do not move the field-update copies earlier.

gpu_velocity_nonlinear!(nl_tor_r, nl_tor_i, nl_pol_r, nl_pol_i, tor_r, tor_i, pol_r, pol_i,
                        config, d1, d2, lfac, rinv, rinv2, rscale, sinθ, cosθ, E, lmax, bw;
                        T_phys=nothing, thermal_factor=zero(eltype(tor_r)), r_vec=nothing,
                        C_phys=nothing, comp_factor=zero(eltype(tor_r)),
                        J_r=nothing, J_θ=nothing, J_φ=nothing,
                        B_r=nothing, B_θ=nothing, B_φ=nothing, lorentz_coeff=zero(eltype(tor_r))) -> nothing

Velocity nonlinear term: nl = analyze( E·(u×ω) − ẑ×u [+ buoyancy + Lorentz] ).

tor/pol are the velocity toroidal/poloidal spectral coefficients; nl_tor/nl_pol receive the toroidal/poloidal nonlinear spectral output. d1/d2 are radial derivative operators, lfac=l(l+1), rinv=1/r, rinv2=1/r², and rscale=1/r² is the radial-synthesis scale forwarded to gpu_vector_spectral_to_physical! (Stage-2 u_r = l(l+1)·P·rscale) — it is LOAD-BEARING, not optional. sinθ/cosθ are the Coriolis grid factors, and E is the Ekman number.

Optional coupled forcing (Phase 5i) accumulated between the Coriolis and the analyze, matching the CPU accumulation order thermal → compositional → Lorentz:

• Thermal buoyancy: supply T_phys (physical (nlat,nlon,nr)), thermal_factor (e.g. (Pm/Pr)·Ra), and r_vec (radial values, length nr). Adds thermal_factor·r·T_phys to the radial component of the advection accumulator.
• Compositional buoyancy: supply C_phys and comp_factor (e.g. (Pm/Sc)·Ra_C). Adds comp_factor·r·C_phys to the radial component.
• Lorentz force: supply J_r/J_θ/J_φ (current density) and B_r/B_θ/B_φ (magnetic field) in physical space, plus lorentz_coeff (e.g. 1/Pm). Adds lorentz_coeff·(J×B) to all three components.

With no keywords (the defaults), the function is exactly the Phase-5g velocity-only path. All arrays (including r_vec, sinθ, cosθ) must live on the same backend (Array or CuArray); mixing host and device arrays errors at broadcast time. (Scratch is pooled via GPUWorkspace when ws is supplied.)

gpu_velocity_poloidal_influence_correction!(x_r, x_i, Gre_b, invG_b) -> nothing

Apply the velocity-poloidal endpoint influence correction in-place to the real (x_r) and imaginary (x_i) parts of a dense (nl,nm,nr) spectral field, using the batched operators from gpu_pack_influence. Backend (CPU/CUDA) is inferred from x_r; x_r, x_i, Gre_b, invG_b must all be on the same backend.

gpu_vr_scale!(vr_alm_r, vr_alm_i, pol_r, pol_i, lfac, rscale) -> nothing

Scale the (split-complex) poloidal coefficients into the vr source coefficients: `vralm[l,m,r] = pol[l,m,r] · lfac[l] · rscale[r].lfac[l+1]=l(l+1)(lengthlmax+1);rscaleis1/r²(the Stage-2 solenoidal convention u_r = l(l+1)·P/r²), lengthnr.lfac/rscalemust reside on the same backend as the coefficient arrays — mixing host and device arrays errors at broadcast time (useon_architecture`).

index_to_lm_shtnskit(idx, lmax, mmax) -> (l, m)

Convert linear spectral index to spherical harmonic degree (l) and order (m).

Index Ordering

The linear index follows the SHTnsKit m-major convention (m varies slowest, l fastest):

• idx=1: (l=0, m=0)
• idx=2: (l=1, m=0)
• idx=3: (l=2, m=0)
• ...
• idx=lmax+1: (l=lmax, m=0)
• idx=lmax+2: (l=1, m=1)
• idx=lmax+3: (l=2, m=1)
• ...

Performance Note

This function uses a linear search. For performance-critical code with many lookups, use index_to_lm_fast with precomputed lookup tables from SHTnsKitConfig.

infer_pencils_from_transpose_name(name::Symbol, plan) -> (src_pencil, dest_pencil)

Infer source and destination pencils from transpose operation name and plan.

GeoDynamo.initialize_fields!(state)

Public entry point for solver field initialization.

This forwards to initialize_solver_fields! so external callers can initialize all enabled fields without depending on the solver-local helper name.

initialize_solver_fields!(state)

Initialize all enabled field families and mark the solver state ready to step.

This is the solver-local implementation behind GeoDynamo.initialize_fields!.

inner_core_alpha(a::InnerCoreAdmittance, l::Int) -> T

ICB admittance α_l for harmonic degree l. Errors if l was not stored.

inner_core_history_flux(a::InnerCoreAdmittance, l, S_old) -> T

ICB radial-derivative contribution φ0 from the inner-core CNAB2 history with a ZERO ICB value: solve M_ic y = b_ic with homogeneous boundary rows (y(0)=0, y(ri)=0) and return φ0 = d1_top · y. Pairs with the admittance α_l to form the outer-core inner Robin row (∂/∂r − α_l) S = φ0. For a zero history, φ0 = 0.

GeoDynamo.install_erk2_cache_bundle!(target, bundle)

Install cache entries from a loaded bundle into a target cache dictionary.

GeoDynamo.install_erk2_cache_bundle!(target::Dict{Symbol, ERK2StageCache{T}}, bundle)

Typed cache-bundle installer used by solver-local cache dictionaries.

integrate_solver_cb3_step!(state)

Advance one complete Cavaglieri-Bewley/Williamson 2N-storage IMEX-RK3 step. solver_step! computes the first nonlinear pass and topography correction before entering this function; substages 2 and 3 recompute them here.

integrate_solver_erk2_step!(state)

Run one full ERK2 timestep for all active solver fields.

The routine builds/reuses field caches, prepares provisional stages, recomputes nonlinear terms at those stages, finalizes each field, applies the velocity-poloidal influence correction, and restores nonlinear histories.

GeoDynamo.load_erk2_cache_bundle!(target, path)

Load a cache bundle from disk, install it into target, and return metadata.

GeoDynamo.load_erk2_cache_bundle(path)

Load ERK2 cache bundle data and metadata from a JLD2 file.

load_parameters([config_file])

Load solver parameters from a file. With no file, loads config/default_params.jl when present, otherwise uses SolverParameters().

load_parameters_from_file(config_file)

Load a Julia parameter file containing assignments that match fieldnames(SolverParameters).

make_subcomms(comm, pencil_r) -> (θ_transform_comm, r_transpose_comm)

Split comm into the two sub-communicators the r×θ decomposition needs, deriving the split colors from the ACTUAL distribution of pencil_r (θ-dist / φ-local / r-dist) so the result is correct for ANY process grid and any PencilArrays rank ordering — NOT from an assumed rank = f(θ_ranks, r_ranks) formula (that only holds when θranks==rranks).

• θ_transform_comm: ranks that share the SAME r-slab and SPLIT θ — the group over which the SH transform distributes θ (so theta_phys/dist_* run here, and the per-level θ-mode gather reduces here). Color = this rank's first owned r index.
• r_transpose_comm: ranks that share the SAME θ-slab and SPLIT r — the group aligned with the r↔lm transpose's radial redistribution. Color = first owned θ index.

GeoDynamo.maybe_log_erk2_stage_residual!(label, buffers, step)

Public wrapper for conditional ERK2 stage-residual logging.

maybe_log_solver_erk2_stage_residual!(label, buffers, step)

Log ERK2 stage residual diagnostics when diagnostics are enabled and step matches the configured interval.

modify_for_flux_inner!(spec_real, spec_imag, slot, prescribed_flux,
                      ∂r, domain, r_range)

Modify coefficients near inner boundary to approximate flux condition. Uses low-order extrapolation.

modify_for_flux_outer!(spec_real, spec_imag, slot, prescribed_flux,
                      ∂r, domain, r_range)

Modify coefficients near outer boundary to approximate flux condition. Uses low-order extrapolation.

normalize_influence_function!(G::Vector{T}, domain::RadialDomain, which_boundary::Int) where T

Normalize influence function to have unit flux at the specified boundary.

on_architecture(arch, array)

Move array to the device associated with arch.

optimize_erk2_transforms!(config::SHTnsKitConfig)

Optimize SHTnsKit transforms for ERK2 timestepping with PencilFFTs. This function pre-warms transform plans and optimizes memory layout.

optimize_fft_performance!(config::SHTnsKitConfig)

Optimize FFT performance by warming up FFTW plans and checking efficiency.

optimize_magnetic_memory_layout!(ℬ::SHTnsMagneticFields{T}) where T

Optimize memory layout for better cache performance using pencil topology

optimize_process_topology(nprocs::Int, dims::Tuple{Int,Int,Int}, spectral_dims=nothing)

Find optimal 2D process grid for given number of processes and problem dimensions. Minimizes communication volume.

optimize_process_topology_shtnskit(nprocs, nlat, nlon, lmax, mmax) -> Tuple{Int,Int}

Find optimal 2D MPI process grid for theta-phi parallelization.

The Optimization Problem

Given nprocs MPI processes, we need to factor it as nprocs = p_theta × p_phi such that:

1. Each process gets at least 2 grid points in each direction
2. Load imbalance (due to non-divisibility) is minimized
3. The decomposition is valid (exact factorization)

Algorithm

Iterates through all valid factorizations and scores them by load imbalance. The load imbalance for a dimension is: |gridsize mod processes| / gridsize

Example

For nprocs=12, nlat=64, nlon=128:

• Possible factorizations: (1,12), (2,6), (3,4), (4,3), (6,2), (12,1)
• (4,3) gives nlat/4=16 points/proc in θ, nlon/3≈43 points/proc in φ
• This might be optimal if it minimizes total imbalance

Arguments

• nprocs: Total number of MPI processes
• nlat: Number of latitude points
• nlon: Number of longitude points
• lmax: Maximum spherical harmonic degree
• mmax: Maximum spherical harmonic order

Returns

• (p_theta, p_phi): Optimal process grid dimensions

optimize_velocity_memory_layout!(𝒰::SHTnsVelocityFields{T}) where T

Optimize memory layout for better cache performance using pencil topology

output_comm()

Return the communicator used by the output layer.

The communicator is resolved lazily through get_comm() so importing the I/O code does not force MPI initialization. The result is cached behind a lock because output calls may be reached from different setup paths.

pack_local_spectral_coefficients(real_data, imag_data, field_info)

Convert local spectral storage into the (spectral_mode, r) slab written to NetCDF.

Mapped local storage keeps the 2D spectral-pencil layout localized to this packing step instead of leaking slot-axis assumptions into I/O code.

parse_proc_grid(spec::Union{AbstractString,Nothing}, nprocs::Int) -> (θ_ranks, r_ranks)

Parse an explicit process grid "θxr" (e.g. "4x2"). At nprocs==1 returns (1,1) without requiring spec. At nprocs>1 spec is REQUIRED and must satisfy θranks·rranks==nprocs.

perform_analysis_from_phi_pencil!(phys_phi, spec, config)

Perform analysis from phi-pencil data.

perform_analysis_phi_local!(phys, spec, config)

Perform analysis when physical field is in phi-pencil orientation.

This is the most efficient analysis path because:

1. The phi (longitude) dimension is fully local on each process
2. SHTnsKit's FFT operates entirely in local memory
3. No MPI communication needed during the transform itself

perform_analysis_with_transpose!(phys, spec, config, to_phi_plan)

Perform analysis with transpose to phi-pencil.

perform_synthesis_phi_local!(spec, phys, config)

Perform synthesis when physical field is in phi-pencil orientation.

This is the most efficient synthesis path because:

1. The phi (longitude) dimension is fully local on each process
2. SHTnsKit's FFT operates entirely in local memory
3. No MPI communication needed during the transform itself

Algorithm for each radial level:

1. Extract spectral coefficients into SHTnsKit's (lmax+1, mmax+1) matrix format
2. Call SHTnsKit.synthesis() which does Legendre transform + FFT
3. Store the resulting (nlat, nlon) physical field slice

perform_synthesis_to_phi_pencil!(spec, phys_phi, config)

Core synthesis routine that writes directly to a phi-pencil array.

This is the workhorse function called by other synthesis routines. It assumes the destination array is already in phi-pencil orientation.

perform_synthesis_with_transpose!(spec, phys, config, back_plan)

Perform synthesis when physical field is NOT in phi-pencil orientation.

Strategy

When the target physical field is in a non-phi pencil (e.g., theta or r pencil), we can't directly use SHTnsKit because it requires all longitude points to be local.

Solution:

1. Create a temporary phi-pencil array
2. Perform synthesis to the temporary array (longitude local)
3. Transpose the result to the target pencil orientation

This involves one extra MPI all-to-all communication (the transpose) but ensures the FFT can operate on contiguous local data.

Pivot singularity threshold: abs(pivot) < eps(T) * PIVOT_SINGULARITY_FACTOR

prepare_solver_erk2_field!(buffers, u, nl, cache, config, dt; bc_spec=nothing)

Prepare the first ERK2 stage for one field.

This computes the full-step linear term, the first nonlinear increment, and the half-step provisional state used for recomputing nonlinear terms.

prettysummary(x)

Compact number formatting for nondimensional quantities (model time, Δt, physics parameters): integers print without a trailing .0, Inf and typemax(Int) print as "Inf", everything else uses Julia's shortest round-trip float printing.

prettytime(t)

Format a duration t in seconds as a human-readable string, e.g. "2.341 seconds", "1.500 days", "100 ns". Follows Oceananigans conventions: picks ns/μs/ms/seconds/minutes/hours/days; integer-valued quantities drop the decimals; second/minute/hour/day pluralize (sub-second unit symbols do not). NaN returns "NaN"; negative durations format as "-<formatted>".

print_shtnskit_config_summary(nlat, nlon, nr, lmax, mmax, nlm, nprocs, memory_estimate)

Print configuration summary for SHTnsKit setup.

rank_seed(seed, rank = get_rank())

Per-rank RNG seed offset for random initial conditions. Each MPI rank fills a DISJOINT subset of spectral modes by stepping a shared rand() stream, so a common seed gave the k-th owned mode the SAME random draw on every rank (spurious cross-rank correlation). Offsetting by rank decorrelates the ranks.

Rank 0 keeps the seed unchanged, so single-rank (nprocs==1) runs are bit-for-bit identical to before. nothing passes through (caller skips seeding and uses the per-process default RNG). NOTE: the resulting IC depends on the rank count (disjoint fill from per-rank streams) — it is reproducible for a fixed (seed, rank-count), not across different rank counts.

Reciprocal condition number threshold for matrix function fallback

read_proc_grid(nprocs::Int) -> (θ_ranks, r_ranks)

Read GEODYNAMO_PROC_GRID from the environment and parse it for nprocs ranks.

read_restart!(tracker, restart_dir, restart_time, config[, pencils]; shtns_config=nothing)

Find and read a restart file near restart_time.

All ranks collectively open the selected file. With pencils, each rank reads only its local slices; without pencils, each rank reads full field arrays. The passed tracker is updated from the restart metadata.

rebuild_solver_implicit_matrices!(state, dt)

Rebuild the pre-factored implicit matrix store for a new timestep dt. The implicit operators bake dt in at factorization time, so changing the timestep requires re-factorization. Reuses the backend's grid/config and all non-timestep physical parameters; only dt changes.

reconstruct_inner_core(a::InnerCoreAdmittance, l, g, S_old) -> Vector

Reconstruct the inner-core radial profile at degree l after the outer-core solve: solve M_ic S = b_ic with regularity S(0)=0 and ICB Dirichlet value S(ri)=g. Returns the length-Nic profile. g is the outer-core value at the ICB; S_old is the previous inner-core profile (its history enters b_ic).

GeoDynamo.report_phi2_conditioning(step; interval=100)

Report phi2 conditioning diagnostics on the requested interval.

reset_energy_tracker!()

Clear all accumulated diagnostic energy histories in ENERGY_TRACKER.

GeoDynamo.reset_phi2_monitor!()

Clear accumulated phi2 conditioning diagnostics.

reset_solver_clock!(state; time, step)

Synchronize the public solver clock and the runtime TimestepState.

Use this when initializing or rewinding a solver state so diagnostics, callbacks, and field views all see the same time/step pair.

resolve_output_precision(sym)

Map user/API precision symbols to concrete floating-point types.

Only :float32 and :float64 are accepted; unknown symbols warn and fall back to Float64 so output remains usable instead of failing late in a write path.

resolved_state_radial_grid(state, radial_grid)

Return the true radial collocation nodes for the writer: an explicit radial_grid wins, otherwise the outer-core RadialDomain carried on state.runtime.outer_core_domain is used when present. Falls back to nothing so the fields-only path keeps working without a runtime.

restore_solver_erk2_nonlinear_terms!(state, temp_buffers, vel_tor_buffers, vel_pol_buffers, mag_tor_buffers, mag_pol_buffers, comp_buffers)

Restore nonlinear terms from pre-stage buffers after an ERK2 step.

The staged nonlinear evaluation temporarily overwrites solver nonlinear fields; this restores the accepted-step histories expected by the rest of the solver.

roll_solver_histories!(state, timestepper, magnetic_enabled)

Advance previous nonlinear histories at the end of a completed timestep.

safe_parse_value(value_str, param_dict)

Parse a parameter value without eval.

GeoDynamo.save_erk2_cache_bundle(path, caches; metadata=Dict())

Persist compatible ERK2 stage caches and metadata to a JLD2 file.

save_parameters(params, filename)

Write a SolverParameters object as plain Julia assignments.

Series term convergence: opnorm(term) < eps(T) * SERIES_CONVERGENCE_FACTOR

set_composition_rhs_bc!(rhs_real, rhs_imag, slot, nr; ...)

Set composition boundary RHS values for one local spectral mode.

GeoDynamo.set_erk2_diagnostics_interval!(interval)

Set how often ERK2 stage residual diagnostics are reported.

set_initial_condition!(model::GeodynamoModel, field::Symbol, ic)

Apply initial condition ic to the named field of model.

Arguments

• model – a GeodynamoModel returned by GeodynamoModel(...)
• field – one of :velocity, :temperature, :magnetic, :composition
• ic – an IC descriptor (RandomPerturbation, AnalyticIC, FileIC, or ZeroIC), or — for the scalar fields :temperature and :composition only — a direct value: a Real (uniform), a function (r, θ, φ) -> value (radius, colatitude, longitude), or an AbstractArray{<:Real,3} matching the local physical grid

Examples

set_initial_condition!(model, :temperature, RandomPerturbation(amplitude=0.1, lmax=10))
set_initial_condition!(model, :magnetic,    AnalyticIC(:dipole; amplitude=1.0))
set_initial_condition!(model, :velocity,    FileIC("/path/to/checkpoint.nc"))
set_initial_condition!(model, :composition, ZeroIC())
set_initial_condition!(model, :temperature, (r, θ, φ) -> 1 - r)

set_magnetic_rhs_bc!(rhs_real, rhs_imag, slot, nr;
                      inner_value=0.0, outer_value=0.0)

Set boundary values in the RHS vector for the magnetic field solve. Matches Fortran magsetbcTor / tim_zerobc: sets RHS boundary rows to zero.

For insulating BCs, all boundary values are zero (homogeneous conditions).

set_scalar_rhs_bc!(rhs_real, rhs_imag, slot, nr; inner_value, outer_value, ...)

Set inner and outer radial RHS entries for one local spectral mode.

set_temperature_rhs_bc!(rhs_real, rhs_imag, slot, nr; ...)

Set temperature boundary RHS values for one local spectral mode.

set_velocity_rhs_bc_poloidal!(rhs_real, rhs_imag, slot, nr;
                               inner_value=0.0, outer_value=0.0)

Set boundary values in the RHS vector for the poloidal velocity solve. Matches Fortran: sets RHS boundary rows to zero for homogeneous BCs.

For no-slip: RHS boundary = 0 (∂P/∂r = 0) For stress-free: RHS boundary = 0 (∂²P/∂r² = 0)

set_velocity_rhs_bc_toroidal!(rhs_real, rhs_imag, slot, nr;
                               inner_value=0.0, outer_value=0.0)

Set boundary values in the RHS vector for the toroidal velocity solve. Matches Fortran velsetbcTor: sets RHS boundary rows to zero (or prescribed value).

For no-slip: RHS boundary = 0 (or prescribed velocity, e.g., rotating inner core) For stress-free: RHS boundary = 0 (homogeneous condition ∂T/∂r - T/r = 0)

setup_diagnostic_variables!(ds, diagnostics, config)

Create scalar diagnostic variables for the diagnostics dictionary.

No variables are created when diagnostics are disabled or the dictionary is empty.

setup_dimensions!(ds, field_info, config)

Define the global NetCDF dimensions shared by coordinates and field variables.

This is part of the collective parallel NetCDF setup path; every rank that opened ds must reach this call in the same order.

setup_variables!(ds, field_info, config, available_fields)

Define NetCDF variables for time, coordinates, and available simulation fields.

The writer defines field variables according to config.output_space: physical scalar fields, spectral coefficient fields, or the mixed layout used by default.

solve_banded!(x, lu, b)

Solve the banded linear system lu * x = b in-place.

In-place aliasing x === b is explicitly supported and safe. The forward substitution sweep processes rows in ascending order: at step i, it reads x[j] for j < i (already holding the intermediate result y[j]) and b[i] (not yet overwritten when aliased). The back substitution sweep processes rows in descending order with the same safe access pattern.

WARNING: Do not change the loop ordering without verifying aliasing safety.

solve_scalar_implicit_step!(solution, rhs, matrices; bc_inner, bc_outer, ...)

Solve a matrix-embedded scalar implicit step for each locally owned spectral mode. Boundary vectors are global mode-indexed arrays; missing values default to homogeneous real and imaginary endpoint values.

solve_to_spec_storage!(cfg, sr, si, solve, plan)

Inverse of spec_storage_to_solve!: scatter the plan-oriented solve's owned m-bins back into the field's spectral storage parents sr/si (which use the even-split (mmax+1) m-partition), performing the θ_comm m-axis redistribution.

solver_build_rhs_cb3_stage!(rhs, u, n, nprev, matrices, dt, gamma, zeta, alpha; mass_coeff, work)

Build the SMR/Cavaglieri-Bewley IMEX-RK3 substage RHS

rhs = (mass_coeff/dt)·u + γ·N + ζ·N_prev + α·(L·u)

consumed by the implicit solve against the system matrix (mass_coeff/dt) I − β·L (built per stage with theta = β). L is the diffusivity-scaled linear operator carried in matrices.linear_matrices; the explicit α·L·u term is the companion Crank-Nicolson half and mirrors the CNAB2 (1−θ)·L·u carry-over in solver_build_rhs_cnab2!. Per-mode and radius-coupled only, so assembly needs no inter-rank communication (each mode owns its full radial profile locally).

solver_build_rhs_cnab2!(rhs, uₙ, nₙ, nₙ₋₁, dt, matrices; mass_coeff=1.0)

Build the CNAB2 right-hand side inside the flattened solver layer in src/ so the solver step reads as a self-contained algorithm instead of reaching back into the legacy timestep entry points.

solver_create_dirichlet_bc(T, nr, value=zero(T))

Create a fixed-value endpoint descriptor for ERK2 radial profiles.

solver_create_insulating_inner_bc(T, d1_row, r_inv)

Create the inner insulating magnetic poloidal endpoint descriptor, (∂r − (l+1)/r)P = 0: under Br = λP/r² the interior vacuum solution is P ∝ r^{l+1} (B = −∇Φ, Φ ∝ r^l regular at the origin ⇒ Br ∝ r^{l−1} = λP/r²). Encoded as lsign = −1, fixedcorrection = −rinv (selfcoeff = d1 − l·rinv − rinv). Matches the banded row in create_magnetic_poloidal_matrices.

solver_create_insulating_outer_bc(T, d1_row, r_inv)

Create the outer insulating magnetic poloidal endpoint descriptor, (∂r + l/r)P = 0: under Br = λP/r² the exterior vacuum solution is P ∝ r^{−l} (B = −∇Φ, Φ ∝ r^{−(l+1)} ⇒ Br ∝ r^{−(l+2)} = λP/r²). Encoded as lsign = +1, fixedcorrection = 0 (selfcoeff = d1 + l·rinv). Verified by the classic full-sphere dipole free-decay rate σ = π² (test/ballbesseldecay.jl); matches the banded row in create_magnetic_poloidal_matrices.

solver_create_neumann_bc(T, d1_row, value=zero(T); l0_dirichlet=false)

Create a first-derivative endpoint descriptor.

solver_create_noslip_pol_bc(T, d1_row)

Create the poloidal velocity no-slip endpoint descriptor.

solver_create_regularity_bc(T, d1_row, r_inv; l_offset=1)

Ball-center regularity endpoint: f′(r₁) = (l + loffset)·f(r₁)/r₁. `loffset = 1for poloidal potentials (f ~ r^{l+1});l_offset = 0` for raw-sphtor toroidal scalars and scalar fields (f ~ r^l; l=0 reduces to f′(r₁)=0).

solver_create_stress_free_pol_bc(T, d2_row)

Create the poloidal velocity stress-free endpoint descriptor.

solver_create_stress_free_tor_bc(T, d1_row, r_inv)

Create the toroidal velocity stress-free endpoint descriptor.

solver_eab2_update_krylov_cached!(u, nₙ, nₙ₋₁, alu_map, domain, ν, config, dt; ...)

Solver-local EAB2 update so the new timestep code does not reach through the legacy timestep API for its exponential update.

The solve gathers a full radial profile per spectral mode, applies the exponential linear action plus the phi1 nonlinear correction, then re-applies any endpoint BCs before scattering back to the local pencil storage.

solver_enforce_erk2_bc!(result, bc_side, boundary_idx, l, nr; value_override=nothing)

Enforce one ERK2 endpoint boundary condition on a dense radial profile.

Dirichlet conditions assign the endpoint directly. Derivative-based conditions solve the endpoint value from the stored stencil and optional l-dependent correction while leaving interior values unchanged.

solver_erk2_stage_residual_stats(buffers)

Return global max and L2 norms of the staged nonlinear residual.

solver_get_eab2_alu_cache!(caches, key, ν, T, domain)

Typed EAB2 cache lookup for the new solver path.

solver_phi2_action_krylov(Aop!, A_lu, v, dt; m, tol)

Compute φ₂(dt·A)·v where φ₂(z) = (eᶻ − 1 − z)/z². Uses the recurrence φ₂(z) = (φ₁(z) − 1)/z, so φ₂(dtA)v = (1/dt)·A⁻¹(φ₁(dtA)v − v) — the same banded A⁻¹ solve φ₁ already uses. Needed for second-order ETD Adams-Bashforth (EAB2): the previous-step nonlinear term must be weighted by φ₂, not φ₁.

solver_solve_composition_implicit_step!(solution, rhs, matrices; kwargs...)

Composition-specific wrapper around the scalar implicit solve.

solver_solve_temperature_implicit_step!(solution, rhs, matrices; kwargs...)

Temperature-specific wrapper around the scalar implicit solve.

solver_solve_velocity_implicit_step!(solution, rhs, matrices, component; ...)

Solve the velocity toroidal or poloidal implicit update with embedded boundary conditions.

Toroidal velocity handles the optional inner-core rotation correction for the l=1, m=0 mode; poloidal velocity uses homogeneous endpoint values here and is corrected later by the ERK2 influence operator when needed.

spec_storage_to_solve!(cfg, solve, sr, si, plan)

Fill the plan-oriented solve PencilArray (parent (l_local, nr, nml)) from the field's spectral storage parents sr/si (parent (l_local, m_local, nr) on the (2,1) spec pencil). Performs the θcomm m-axis redistribution. solve columns `1:length(plan.mlocal)are filled (degreeplan.m_local[mi]`); any trailing bins are zeroed.

spectral_mode_counts(θ_ranks, r_ranks, lmax, mmax) -> Vector{Int}

Number of spherical-harmonic modes each rank owns under the spectral pencil's contiguous-block split (m over θ_ranks, l over r_ranks). Because of triangular truncation (valid modes need m ≤ l), equal SLOT blocks give unequal MODE counts, and a square grid can leave a rank with ZERO modes (high-m / low-l corner). Used to warn about idle ranks and load imbalance at grid setup. Pure function (no MPI).

store_coefficients_from_shtnskit!(spec_real, spec_imag, coeffs_matrix, r_local, config)

Convert SHTnsKit coefficient matrix format back to our linear spectral storage.

This is the inverse of extract_coefficients_for_shtnskit!:

• Input: SHTnsKit's (lmax+1) × (mmax+1) coefficient matrix
• Output: Our linear-indexed spectral arrays (real and imaginary parts separate)

Physical Constraint

For real-valued physical fields, the m=0 coefficients must be purely real. This function enforces this by zeroing the imaginary part for m=0 modes.

store_physical_slice_generic!(phys_data, phys_slice, r_local, config)

Generic storage for any pencil orientation.

store_physical_slice_phi_local!(phys_data, phys_slice, r_local, config)

Copy a 2D physical field slice into a 3D array at a specific radial level.

Optimized for Phi-Local Layout

When the physical field is in phi-pencil orientation, the (θ, φ) indices correspond directly to the array's first two dimensions, making this a straightforward copy operation.

Arguments

• phys_data: 3D destination array (θ × φ × r)
• phys_slice: 2D source array from SHTnsKit (θ × φ)
• r_local: Radial index to write to
• config: Configuration for grid dimensions

store_solver_erk2_stage_nonlinear!(buffers, nl)

Store nonlinear terms evaluated at the provisional ERK2 stage.

store_zero_component_generic!(v_component, r_local, config)

Set a component to zero at a given radial level. Used at r=0 (ball geometry) where v_r must be zero for regularity.

sync_gpu_state_to_cpu!(cpu_state, gpu_state) -> cpu_state

Write the GPU device-state back into the CPU SolverState so CPU stepping / diagnostics / output / restart can continue from the GPU-evolved state. Three categories are synced (the inverse of what build_gpu_solver_state reads):

1. spectral fields (velocity tor/pol, temperature, and — when present — magnetic tor/pol, composition), into the slot-packed CPU storage;
2. the CNAB2 prev_nl histories of the same fields;
3. the lagged physical buffers (T_phys/C_phys, B_*/J_*) into the CPU physical fields they were built from.

The gpu_state arrays may be host or device.

sync_spectral_history!(previous::SHTnsSpecField, current::SHTnsSpecField)

Copy spectral coefficients from current into previous. This helper is shared by the solver rewrite path and any remaining legacy timestepping code that keeps one-step nonlinear histories.

synchronize_pencil_data!(field)

Synchronize PencilArray data across MPI processes to ensure consistency.

synchronize_pencil_transforms!(field::SHTnsSpecField{T}) where T

Ensure all pending PencilFFTs operations are completed and data is consistent across processes.

time_to_next_output(tracker, current_time, config)

Return the smallest wait time until the next history or restart event.

This helper is used by scheduling code that wants a conservative next wake-up time without duplicating the tracker cadence logic.

to_sht_spectral(cfg, alm_mat::AbstractMatrix{ComplexF64})

Hand a full (lmax+1, mmax+1) coefficient matrix to the distributed synthesis. The matrix is already the dense form dist_synthesis expects, so this is the identity.

transform_field_and_gradients_to_physical!(𝔽::AbstractScalarField{T}, ws::GradientWorkspace{T}) where T

Transform scalar field and all gradient components to physical space in a single batched operation to minimize communication.

unpack_local_spectral_coefficients(real_data, imag_data, config)

Rebuild local spectral storage from a NetCDF (spectral_mode, r) slab using the configured local spectral-slot mapping.

update_tracker!(tracker, current_time, config, did_output, did_restart)

Advance output and restart counters after a write attempt.

Pass the actual write decisions so history and restart cadence can advance independently when only one output type was produced.

validate_magnetic_configuration(ℬ::SHTnsMagneticFields{T}, config::C) where {T,C<:SHTnsKitConfig}

Validate magnetic field configuration consistency with SHTns setup

validate_mpi_consistency!(field::SHTnsSpecField{T}) where T

Check that spectral field data is valid (no NaN/Inf) across all MPI processes. Returns the field after validation. Warns if any process has invalid data.

validate_output(filename)

Return true when filename can be opened as NetCDF and has the minimum history-file variables required by the output reader helpers.

validate_output_compatibility(field_info, shtns_config)

Check that field metadata and SHTns metadata describe the same grid and spectral layout.

Returns false and emits a warning listing all mismatches; returns true when the layouts are compatible.

validate_parameters(params; strict=false)

Validate a SolverParameters object.

validate_pencil_decomposition(config::SHTnsKitConfig)

Validate that pencil decomposition is optimal for the problem size and MPI configuration.

validate_proc_grid(θ_ranks, r_ranks; nlat, nr, lmax, mmax)

Warn (once, on rank 0) about a process grid that over-decomposes an axis (idle ranks) or yields a large spherical-harmonic mode-load imbalance. Purely advisory — does not change the decomposition.

validate_velocity_configuration(𝒰::SHTnsVelocityFields{T}, config::C) where {T,C<:SHTnsKitConfig}

Validate velocity field configuration consistency with SHTns setup

verify_all_ranks_wrote(output_dir, hist_number; geometry="shell", expected_dims=nothing)

Verify that the single parallel output file exists for a given history number and has correct global dimensions.

with_output_precision(config, T)

Return a copy of config that writes field and scalar output with precision T.

All scheduling, metadata, directory, and layout settings are preserved.

write_coordinate_data!(ds, field_info, config)

Write coordinate arrays. Only rank 0 writes coordinates (they are global/shared).

write_diagnostics!(ds, diagnostics, config)

Write scalar diagnostics into diag_* variables when diagnostics are enabled.

The values are expected to be already globally reduced by compute_diagnostics.

write_field_data!(ds, fields, config, field_info)

Write field data using parallel offset writes. Each rank writes its local pencil slice at the correct global position.

write_fields!(state, tracker, metadata, config, shtns_config, pencils)
write_fields!(fields, tracker, metadata, config, shtns_config, pencils)

All-ranks entry point: synchronizes the output decision via MPI.Bcast, then extracts fields only if output is actually needed.

The state overload accepts any runtime object for which extract_all_fields is defined. The fields overload accepts an already extracted field dictionary. Both overloads must be called by every rank on every timestep to avoid MPI deadlocks in collective output decisions.

write_grid_file!(config, field_info, shtns_config, metadata)

Write a separate grid file containing coordinate and grid information. Written only once by rank 0 at the start of the simulation.

write_restart!(fields, tracker, metadata, config[, pencils]; shtns_config=nothing, geometry=:shell, radius_ratio=0.35)

Write a restart NetCDF file using the same parallel field layout as history output.

The restart file also stores enough TimeTracker state for a resumed run to continue output and restart numbering without clobbering earlier files.

write_time_data!(ds, time, step, config)

Write scalar simulation time and step values.

Only rank 0 writes these shared scalar variables; field arrays are handled by the distributed write path.

zero_scalar_work_arrays!(𝔽::AbstractScalarField{T}) where T

Efficiently zero all work arrays for a scalar field.