API reference

The stable public facade is dkx.api (re-exported at the package top level). The canonical modules that implement the physics and solver stack are indexed below, with links to the pages that document them in depth and to Source-code map for the full source catalogue.

Public facade

Stable public data contracts for high-level DKX workflows.

The classes in this module describe orchestration boundaries: user inputs, geometry/grid/operator summaries, solver metadata, output schemas, and benchmark reports. They intentionally avoid importing JAX so they stay cheap to import from the CLI, documentation, tests, and downstream workflow code.

class dkx.api.BenchmarkReport(case: str, backend: str, runtime_s: float, peak_memory_mb: float, status: str, metadata: Mapping[str, ~typing.Any]=<factory>): Runtime, memory, and status summary for one benchmark case.

class dkx.api.GeometryState(kind: str, source_path: Path | None = None, radial_coordinate: float | None = None, metadata: Mapping[str, ~typing.Any]=<factory>): Geometry contract exposed to problem setup and validation layers.

class dkx.api.GridState(n_theta: int, n_zeta: int, n_xi: int, n_x: int, n_species: int = 1, metadata: Mapping[str, ~typing.Any]=<factory>): Phase-space grid sizes and layout metadata.

class dkx.api.OperatorState(rhs_mode: int, size: int, collision_model: str, include_phi1: bool = False, matrix_free: bool = True, metadata: Mapping[str, ~typing.Any]=<factory>): Operator summary shared by problem, solver, and diagnostics layers.

class dkx.api.OutputSchema(format: str, version: str, keys: Sequence[str] = (), path: Path | None = None, metadata: Mapping[str, ~typing.Any]=<factory>): Versioned output-file schema and field-key contract.

class dkx.api.PreconditionerState(kind: str, differentiable: bool, device_safe: bool, setup_memory_mb: float | None = None, metadata: Mapping[str, ~typing.Any]=<factory>): Selected preconditioner capability and setup metadata.

class dkx.api.SolveInputs(input_path: Path | None = None, wout_path: Path | None = None, output_path: Path | None = None, backend: str | None = None, requires_autodiff: bool = False, options: Mapping[str, ~typing.Any]=<factory>): Normalized high-level solve request passed across API/CLI boundaries.

class dkx.api.SolverResult(residual_norm: float, converged: bool, iterations: int | None = None, runtime_s: float | None = None, solution: Any = None, metadata: Mapping[str, ~typing.Any]=<factory>): Backend-neutral solver summary for API-level callers.

class dkx.api.TransportResult(transport_matrix: Any = None, particle_flux: Any = None, heat_flux: Any = None, bootstrap_current: Any = None, solver: SolverResult | None = None, metadata: Mapping[str, ~typing.Any]=<factory>): Transport-output summary independent of a specific file format.

Batched E_r scan on one geometry (stable public facade).

Routes to dkx.batch.batched_er_scan(): one jax.vmap batched solve over a vector of radial-electric-field values sharing a single geometry, returning batched states, moments, and the radial current J_r per E_r. Auto-chunked to a memory-budgeted batch size; differentiable and jit-safe. The heavy JAX/batch stack is imported lazily so dkx.api stays cheap to import.

Parameters:

request – a prepared dkx.er.ErProblem, or a deck (SolveInputs / namelist path) that is prepared with dkx.er.prepare() (RHSMode=1, inputRadialCoordinate=4 Er knob).
er_values – the E_r scan values, shape (batch,).
er_bracket – optional bracket / initial E_r forwarded to dkx.er.prepare() when request is a deck.
er_initial – optional bracket / initial E_r forwarded to dkx.er.prepare() when request is a deck.
solve_method – forwarded to the solve.
tol – forwarded to the solve.
differentiable – differentiable implicit solves (for jax.grad).
max_batch – optional memory-budgeting overrides.
memory_budget_gb – optional memory-budgeting overrides.

Returns:

A dkx.batch.BatchedSolveResult (radial_current populated).

dkx.api.batched_surface_scan(surfaces: Sequence[Any], *, solve_method: str = 'auto', tol: float = 1e-10, differentiable: bool = False, max_batch: int | None = None, memory_budget_gb: float | None = None) → Any

Batched solve over a batch of flux surfaces (stable public facade).

Routes to dkx.batch.batched_surface_scan(): a jax.vmap batched solve over a sequence of flux-surface operators that share discretization (grids/derivative matrices/layout) but differ in geometry, species, collision, and drive leaves. Auto-chunked to a memory-budgeted batch size; differentiable and jit-safe.

Parameters:

surfaces – a sequence whose entries are each a built dkx.drift_kinetic.KineticOperator, or a deck (SolveInputs / namelist path) built into one per surface.
solve_method – forwarded to the solve.
tol – forwarded to the solve.
differentiable – forwarded to the solve.
max_batch – optional memory-budgeting overrides.
memory_budget_gb – optional memory-budgeting overrides.

Returns:

A dkx.batch.BatchedSolveResult with batched states and moments.

dkx.api.bounce_averaged_transport(geometry: Any, *, r_eff: float | None = None, grad_psi_avg: float | None = None, n_field_periods: int = 160, points_per_period: int = 48, n_pitch: int = 128, n_quad: int = 14, max_wells: int = 224, n_field_lines: int = 1, m_keep: int | None = None, n_keep: int | None = None) → Any

Bounce-averaged 1/nu effective-ripple transport (stable public facade).

Routes to dkx.bounce_averaged.bounce_averaged_transport(): the differentiable, radially-local, bounce-averaged surrogate for the dominant low-collisionality (1/nu-regime) neoclassical radial transport of a flux surface – the effective ripple epsilon_eff and the trapped-particle bounce integrals of the radial magnetic drift – computed from the |B| Boozer spectrum of geometry (J. L. Velasco et al., J. Comput. Phys. 418, 109512 (2020); Nucl. Fusion 61, 116059 (2021); spectrally-accurate differentiable bounce points as in arXiv:2412.01724). Pure JAX, jit/vmap safe, and differentiable in the Boozer amplitudes when geometry is built with FluxSurfaceGeometry.from_fourier(). The heavy JAX stack is imported lazily so dkx.api stays cheap to import.

Parameters:

geometry – a dkx.magnetic_geometry.FluxSurfaceGeometry.
r_eff – the <|grad psi|> normalization of epsilon_eff (large-aspect-ratio B_0 r_eff or an explicit average); the |grad psi|-free gamma_c core needs neither.
grad_psi_avg – the <|grad psi|> normalization of epsilon_eff (large-aspect-ratio B_0 r_eff or an explicit average); the |grad psi|-free gamma_c core needs neither.
n_field_periods – quadrature/bandwidth controls (see the owning module).
points_per_period – quadrature/bandwidth controls (see the owning module).
n_pitch – quadrature/bandwidth controls (see the owning module).
n_quad – quadrature/bandwidth controls (see the owning module).
max_wells – quadrature/bandwidth controls (see the owning module).
n_field_lines – quadrature/bandwidth controls (see the owning module).
m_keep – quadrature/bandwidth controls (see the owning module).
n_keep – quadrature/bandwidth controls (see the owning module).

Returns:

A dkx.bounce_averaged.BounceAveragedTransport (.epsilon_eff, .epsilon_eff_32, .gamma_c, …).

dkx.api.momentum_corrected_bootstrap(database: Any, *, z_s: Any, m_hats: Any, n_hats: Any, t_hats: Any, nu_n: Any, dn_hat_dpsi_hat: Any, dt_hat_dpsi_hat: Any, dphi_hat_dpsi_hat: Any = 0.0, e_par_b: Any = 0.0, uncorrected_flows: Any = None, x: Any = None, x_weights: Any = None, n_x: int = 64, x_max: float = 5.0) → Any

Momentum-corrected bootstrap current (stable public facade).

Routes to dkx.momentum_correction.momentum_corrected_bootstrap(), the Sugama-Nishimura moment-method parallel-momentum correction on the monoenergetic transport coefficients (H. Sugama and S. Nishimura, Phys. Plasmas 9, 4637 (2002); 15, 042502 (2008); H. Maassberg, C. D. Beidler, and Y. Turkin, Phys. Plasmas 16, 072504 (2009)). The heavy JAX stack is imported lazily so dkx.api stays cheap to import.

Parameters:

database – a dkx.monoenergetic.MonoenergeticDatabase (from run_monoenergetic_database()).
z_s – species parameters, shape (S,).
m_hats – species parameters, shape (S,).
n_hats – species parameters, shape (S,).
t_hats – species parameters, shape (S,).
nu_n – deck normalized collisionality.
dn_hat_dpsi_hat – radial gradients dn/dpsiHat, dT/dpsiHat per species (shape (S,)).
dt_hat_dpsi_hat – radial gradients dn/dpsiHat, dT/dpsiHat per species (shape (S,)).
dphi_hat_dpsi_hat – radial-electric-field and inductive parallel-field drives (default 0).
e_par_b – radial-electric-field and inductive parallel-field drives (default 0).
uncorrected_flows – optional (S,) override for the uncorrected parallel-flow moments (e.g. from a kinetic solve).
x – speed quadrature controls.
x_weights – speed quadrature controls.
n_x – speed quadrature controls.
x_max – speed quadrature controls.

Returns:

A dkx.momentum_correction.MomentumCorrectionResult (.corrected_bootstrap, .uncorrected_bootstrap, .delta_bootstrap, .corrected_flows, …).

dkx.api.read_output(path: str | Path) → dict[str, Any]: Read an HDF5, NetCDF, or NPZ dkx output file.

dkx.api.run_ambipolar_brent(request: SolveInputs | str | Path, *, work_dir: str | Path | None = None, er_min: float, er_max: float, er_initial: float = 0.0, max_evaluations: int = 20, current_tolerance: float = 1e-10, step_tolerance: float = 1e-08, solve_method: str = 'auto', differentiable: bool | None = None, reuse_output_geometry_cache: bool = True, reuse_solver_state: bool = True, emit: Any = None, **kwargs: Any) → Any

Run the canonical-stack Brent ambipolar E_r solve (stable public facade).

Routes to dkx.er.find_ambipolar_er() (the canonical inputs -> drift_kinetic -> solve -> moments slice, replacing the legacy in-process Brent owner). work_dir, step_tolerance, reuse_output_geometry_cache and reuse_solver_state are accepted for backwards compatibility; warm starts / recycling are threaded internally.

Returns:: A dkx.er.AmbipolarResult (.er is the ambipolar E_r; .roots lists every classified root).

dkx.api.run_monoenergetic_database(request: SolveInputs | str | Path, nu_prime_grid: Sequence[float], e_star_grid: Sequence[float] = (0.0,), *, output_path: str | Path | None = None, solve_method: str = 'auto', tol: float = 1e-10, emit: Any = None) → Any

Scan (nuPrime, EStar) monoenergetic coefficients (stable public facade).

Routes to dkx.monoenergetic.monoenergetic_database() (the canonical RHSMode=3 database scan) and optionally writes the compact .npz database via dkx.monoenergetic.save_database(). The heavy stack is imported lazily so dkx.api stays cheap to import.

Returns:: The dkx.monoenergetic.MonoenergeticDatabase.

dkx.api.write_output(request: SolveInputs | str | Path, output_path: str | Path | None = None, **kwargs: Any) → Path

Run dkx from Python and write an output file.

request may be a SolveInputs object or an input namelist path. Routes to dkx.run.run_from_namelist() (the canonical RHSMode dispatch behind dkx write-output). The implementation imports the heavy run stack lazily so importing dkx.api stays cheap for docs, CLI validation, and downstream workflow planning.

High-level runners live in dkx.run (run_profile, run_transport_matrix) and the CLI in dkx.cli; see Usage.

Canonical modules

Module	Role	Reference
`dkx.run`	`run_profile` / `run_transport_matrix` orchestration and console flow	Usage
`dkx.inputs` / `dkx.namelist`	Typed `SfincsInput` parser and raw namelist reader	Inputs (namelist) reference
`dkx.magnetic_geometry`	`FluxSurfaceGeometry` and all geometry schemes	Geometry models and loading
`dkx.species`	Charges, masses, profiles, deflection frequencies	Physics reference: the radially local drift-kinetic model
`dkx.phase_space`	Legendre coupling, Landreman–Ernst speed grid, `Nxi_for_x` ramp	Numerics and algorithms
`dkx.drift_kinetic`	`KineticOperator` — the consolidated v3 drift-kinetic operator	Physics reference: the radially local drift-kinetic model, Drift-kinetic equation and system of equations
`dkx.collisions`	Pitch-angle scattering and full Fokker–Planck (Rosenbluth) operators	Physics reference: the radially local drift-kinetic model
`dkx.solve`	Three-tier auto solver (block-tridiagonal, recycled Krylov, host direct)	Numerics and algorithms
`dkx.moments`	Velocity-space moments, fluxes, `FSABjHat`, transport matrix	Outputs (HDF5, NetCDF4, and NPZ), Physics reference: the radially local drift-kinetic model
`dkx.phi1`	Nonlinear \(\Phi_1\) / quasineutrality Newton solve	Physics reference: the radially local drift-kinetic model
`dkx.er`	Ambipolar radial-electric-field root solve (Brent + differentiable)	Physics reference: the radially local drift-kinetic model
`dkx.writer` / `dkx.io`	`sfincsOutput.h5` / `.nc` / `.npz` writer and reader	Outputs (HDF5, NetCDF4, and NPZ)
`dkx.sensitivity`	Implicit-differentiation observable derivatives (RHSMode 4/5 spine)	Fortran v3 And dkx Feature Matrix
`dkx.compare`	`compare-h5` parity tooling against Fortran fixtures	Parity status

Namelist parsing and reader helpers

Minimal parser for SFINCS Fortran namelist files.

Goals: - No third-party dependency (no f90nml). - Parse the SFINCS v3 example input.namelist files in sfincs/fortran/version3/examples.

This is not a complete Fortran namelist implementation.

class dkx.namelist.Namelist(groups: 'Dict[str, Dict[str, Value]]', indexed: 'Dict[str, Dict[str, Dict[Tuple[int, ...], Scalar]]]', source_path: 'Path | None' = None, source_text: 'str | None' = None)

dkx.namelist.parse_sfincs_input_text(text: str, *, source_path: str | Path | None = None) → Namelist

Parse SFINCS namelist text into groups.

This helper is useful for compact regression fixtures that embed small Fortran-v3 decks without keeping a directory of one-file inputs.

dkx.namelist.read_sfincs_input(path: str | Path) → Namelist: Parse a SFINCS input.namelist file into groups.

Reduced-model and analysis modules

Module	Role	Reference
`dkx.monoenergetic`	Monoenergetic-database mode: `(nuPrime, EStar)` scans and energy convolution to thermal `L_ij`	Reduced-model capabilities
`dkx.variational`	Variational upper/lower bounds on the monoenergetic \(D_{11}\) (convergence certificate)	Reduced-model capabilities
`dkx.shaing_callen`	Collisionless-limit bootstrap coefficient with an analytic axisymmetric cross-check	Reduced-model capabilities

The differentiable solve, the implicit adjoint, and the vmex -> booz_xform_jax -> dkx chain are documented in Differentiability; the full source catalogue is Source-code map.