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.SfincsInput(general: GeneralParams = <factory>, geometry: GeometryParams = <factory>, species: SpeciesParams = <factory>, physics: PhysicsParams = <factory>, resolution: ResolutionParams = <factory>, other: OtherNumericalParams = <factory>, preconditioner: PreconditionerOptions = <factory>, export_f: Mapping[str, str | bool | int | float | ~typing.List[str | bool | int | float]]=<factory>, raw: RawNamelist | None = None, warnings: Tuple[str, ...]=())

Typed SFINCS input, grouped by namelist section.

export_f and anything not (yet) typed remain reachable through raw (the full parsed namelist) — the .raw fallback contract.

Build one from a file with load_sfincs_input(), or programmatically with from_params() using the flat Fortran parameter names (SfincsInput.from_params(Ntheta=17, geometryScheme=11, ...)); the nested section dataclasses use snake_case Python names (n_theta) whose Fortran spelling is recorded in each field’s nml metadata. Serialize back to Fortran-namelist text with to_namelist() or write().

classmethod from_params(*, validate: bool = True, **params: Any) → SfincsInput

Build a typed input from flat Fortran-named parameters.

Parameter names are the Fortran input.namelist names (Ntheta, geometryScheme, equilibriumFile, Zs, …), matched case-insensitively and routed to the owning namelist section automatically. Species arrays accept lists or tuples. Unknown names raise ValueError (the .raw fallback of parsed decks has no programmatic equivalent).

Parameters:

validate – run the validateInput.F90-equivalent checks and the RHSMode=3 hard overrides (default), as load_sfincs_input() does for files.
**params – flat Fortran-named values, e.g. SfincsInput.from_params(geometryScheme=1, Ntheta=15, Zs=[1.0]).

Returns:

The typed (and, by default, validated) SfincsInput with raw=None; the run drivers synthesize the namelist via to_namelist() when needed.

to_namelist(*, include_defaults: bool = False) → str

Serialize to SFINCS-style Fortran input.namelist text.

The output is the parser’s inverse: load_sfincs_input of the written text reproduces every typed section field, the export_f group, untyped keys retained in raw (legacy aliases such as JGboozer_file), and the rank-2 boozer_bmnc(m,n) spectra of geometryScheme=13 decks. The formatting (&group … /, .true./.false., quoted strings, space-separated arrays) is the Fortran-namelist subset both this parser and the SFINCS v3 Fortran reader accept.

Parameters:: include_defaults – write every typed field; the default writes a compact deck with only the fields that differ from the Fortran defaults (globalVariables.F90).

write(path: str | Path, *, include_defaults: bool = False) → Path: Write to_namelist() text to path and return the Path.

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.SolverOptions(method: str = 'auto', tol: float = 1e-10, atol: float = 0.0, restart: int = 30, recycle_dim: int = 8, max_restarts: int = 200, differentiable: bool = False, use_preconditioner: bool = True, device: Any = None, memory_budget_gb: float | None = None, cores: int | None = None)

Typed tuning knobs for dkx.solve.solve(), threaded by the run drivers.

Pass to dkx.run.run_profile(), dkx.run.run_transport_matrix(), or dkx.run.run_from_namelist() as solver=SolverOptions(...); when given, it supersedes the drivers’ quick solve_method/tol arguments and is expanded to dkx.solve.solve() keywords via solve_kwargs(). Environment variables keep working as overrides for knobs left at None (memory_budget_gb=None reads DKX_TIER1_MEMORY_BUDGET_GB inside the solve).

method

"auto" | "block_tridiagonal" | "gmres" | "direct" (the three-tier policy of dkx.solve.solve()).

Type:: str

tol

relative residual tolerance (per RHS column).

Type:: float

atol

absolute residual floor.

Type:: float

restart

FGMRES cycle size m (tier 2).

Type:: int

recycle_dim

GCROT recycle directions k (tier 2).

Type:: int

max_restarts

tier-2 outer-cycle cap (the tier-3 trigger in auto).

Type:: int

differentiable

wrap the solution in an implicit-function-theorem linear_solve so jax.grad flows through (tiers 1/2).

Type:: bool

use_preconditioner

tier-2 coarse-operator preconditioner on/off.

Type:: bool

device

JAX device for the solve (a platform string such as "cpu" or "gpu", or a concrete jax.Device); None keeps the operator’s placement.

Type:: Any

memory_budget_gb

budget above which method="auto" prefers the memory-lean truncated tier-1 kernel over the full-band factorization; None reads DKX_TIER1_MEMORY_BUDGET_GB, else the solve’s default applies.

Type:: float | None

cores

host CPU threadpool width. XLA sizes its threadpool once, before the first JAX device use, so a value stored here CANNOT change a process whose JAX backend is already initialized; it is carried for provenance and planning only. To actually pin threads, set the DKX_CORES environment variable or the CLI --cores flag before import dkx (see docs/parallelism.rst); solve_kwargs() deliberately excludes this field.

Type:: int | None

solve_kwargs() → dict[str, Any]: Keyword arguments for dkx.solve.solve() (cores excluded).

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.

Run dkx from Python and write an output file.

request may be a SolveInputs object, an input namelist path, or an in-memory dkx.inputs.SfincsInput (no file is read or written for the input; the output’s input.namelist provenance dataset stores to_namelist() text when the input was built programmatically). 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.