DIRKStep

DIRKStep provides a diagonally implicit Runge–Kutta integrator that solves stage systems with the cached Newton–Krylov helpers supplied by cubie.integrators.matrix_free_solvers. The factory consumes DIRKTableau instances, exposing L-stable SDIRK and ESDIRK schemes for stiff problems while preserving adaptive error control through embedded weights.

Defaults

algorithm="dirk" integrates with the l_stable_dirk_3 tableau (L-stable, stiffly accurate, three stages, order 3). That tableau has no embedded error estimate, so the default is fixed-step control; choosing a DIRK tableau that provides an estimate enables the family’s Gustafsson predictive defaults (step-size growth clamped to 0.2–8.0×). Each stage runs one Newton–Krylov solve with the shared defaults listed in Default settings. The implicit-midpoint, SDIRK, and Hairer–Wanner L-stable schemes on the DIRK tableau registry page resolve to this class.

class cubie.integrators.algorithms.DIRKStep(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], n_states: int, dxdt_fn: Callable | None = None, observables_fn: Callable | None = None, drivers_fn: Callable | None = None, get_solver_helper_fn: Callable | None = None, tableau: DIRKTableau = DIRKTableau(a=((0.43586652150845895, 0.0, 0.0), (0.2820667392457705, 0.43586652150845895, 0.0), (1.2084966491760099, -0.6443631706844688, 0.43586652150845895)), b=(1.2084966491760099, -0.6443631706844688, 0.43586652150845895), c=(0.43586652150845895, 0.7179332607542295, 1.0), order=3, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=0.85, dense_prediction_ratio_float64=1.07, defaults={}), n_drivers: int = 0, attempt_dense_prediction: bool = True, **kwargs)[source]

Bases: ODEImplicitStep

Diagonally implicit Runge–Kutta step with an embedded error estimate.

__init__(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], n_states: int, dxdt_fn: Callable | None = None, observables_fn: Callable | None = None, drivers_fn: Callable | None = None, get_solver_helper_fn: Callable | None = None, tableau: DIRKTableau = DIRKTableau(a=((0.43586652150845895, 0.0, 0.0), (0.2820667392457705, 0.43586652150845895, 0.0), (1.2084966491760099, -0.6443631706844688, 0.43586652150845895)), b=(1.2084966491760099, -0.6443631706844688, 0.43586652150845895), c=(0.43586652150845895, 0.7179332607542295, 1.0), order=3, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=0.85, dense_prediction_ratio_float64=1.07, defaults={}), n_drivers: int = 0, attempt_dense_prediction: bool = True, **kwargs) None[source]

Initialise the DIRK step configuration.

This constructor creates a DIRK step object and automatically selects appropriate default step controller settings based on whether the tableau has an embedded error estimate. Tableaus with error estimates default to adaptive stepping with order-dependent PI controller defaults, while errorless tableaus default to fixed stepping.

Parameters:
  • precision – Floating-point precision for CUDA computations.

  • n_states – Number of state variables in the ODE system.

  • dxdt_fn – Device function for evaluating f(t, y) right-hand side.

  • observables_fn – Device function computing system observables.

  • drivers_fn – Optional device function evaluating drivers at arbitrary times.

  • get_solver_helper_fn – Factory function returning solver helper for Jacobian operations.

  • tableau – DIRK tableau describing the coefficients. Defaults to DEFAULT_DIRK_TABLEAU.

  • n_drivers – Number of driver variables in the system.

  • attempt_dense_prediction – Request dense stage prediction; ignored when the tableau does not meet the transform’s preconditions.

  • **kwargs – Optional parameters passed to config classes. See DIRKStepConfig, ImplicitStepConfig, and solver config classes for available parameters. None values are ignored.

Notes

The step controller defaults are selected dynamically:

  • If tableau.has_error_estimate is True: Uses DIRK_ADAPTIVE_DEFAULTS (order-dependent PI controller defaults)

  • If tableau.has_error_estimate is False: Uses DIRK_FIXED_DEFAULTS (fixed-step controller)

This automatic selection prevents incompatible configurations where an adaptive controller is paired with an errorless tableau.

_build_error_solver() None[source]

Construct the width-n smoothing solver from live settings.

property accumulator_elements: int

Elements of the explicit-stage accumulator.

build_implicit_helpers() None[source]

Request the helpers and push the solver chain’s products.

build_step(dxdt_fn: Callable, observables_fn: Callable, drivers_fn: Callable | None, solver_function: Callable, numba_precision: type, n: int, n_drivers: int) StepCache[source]

Compile the DIRK device step.

default_tableau = DIRKTableau(a=((0.43586652150845895, 0.0, 0.0), (0.2820667392457705, 0.43586652150845895, 0.0), (1.2084966491760099, -0.6443631706844688, 0.43586652150845895)), b=(1.2084966491760099, -0.6443631706844688, 0.43586652150845895), c=(0.43586652150845895, 0.7179332607542295, 1.0), order=3, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=0.85, dense_prediction_ratio_float64=1.07, defaults={})

Tableau a bare family alias builds on; None for fixed schemes.

classmethod family_defaults(tableau=None) AlgorithmDefaults[source]

Adaptive or fixed defaults by the tableau’s error estimate.

property has_error_estimate: bool

Return True when the tableau supplies an error estimate.

property is_implicit: bool

Return True because the method solves nonlinear systems.

property is_multistage: bool

Return True as the method has multiple stages.

property newton_solves_per_step: int

one per implicit stage.

Type:

Newton solves one step runs

property optimisation_candidates: Tuple[Dict[str, Any], ...]

Newton unrolling crossed with accumulator placement, plus rolled other_small at rolled Newton with a local accumulator.

property order: int

Return the classical order of accuracy.

performance_defaults(hardware: Any = None) Dict[str, Any][source]

Share a spilling Krylov accumulator while the GPU stays at least half full.

register_buffers() None[source]

Register buffers according to locations in compile settings.

property threads_per_step: int

Return the number of CUDA threads that advance one state.

class cubie.integrators.algorithms.generic_dirk.DIRKStepConfig(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], n_states: int = 1, n_drivers: int = 0, is_adaptive: bool = True, dxdt_fn: Callable | None = None, observables_fn: Callable | None = None, drivers_fn: Callable | None = None, get_solver_helper_fn: Callable | None = None, operator_beta: float = 1.0, operator_gamma: float = 1.0, preconditioner_order: int | None = None, preconditioner_type: str = 'jacobi', use_smoothed_error: bool = False, inexact_newton: bool = False, prefactored: bool = True, cached_auxiliaries_location: str = 'local', newton_nonlinear_solver_fn: Callable | None = None, krylov_linear_solver_fn: Callable | None = None, prepare_jacobian_fn: Callable | None = None, error_linear_solver_fn: Callable | None = None, helper_operation_counts: OperationCounts = NOTHING, tableau: DIRKTableau = DIRKTableau(a=((0.43586652150845895, 0.0, 0.0), (0.2820667392457705, 0.43586652150845895, 0.0), (1.2084966491760099, -0.6443631706844688, 0.43586652150845895)), b=(1.2084966491760099, -0.6443631706844688, 0.43586652150845895), c=(0.43586652150845895, 0.7179332607542295, 1.0), order=3, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=0.85, dense_prediction_ratio_float64=1.07, defaults={}), attempt_dense_prediction: bool = True, predictor_fn: Callable | None = None, stage_increment_location: str = 'local', stage_increment_history_location: str = 'local', previous_step_size_location: str = 'local', stage_base_location: str = 'local', accumulator_location: str = 'local', stage_rhs_location: str = 'local', apply_mass_fn: Callable | None = None, inverse_mass_dxdt_fn: Callable | None = None, *, jit_flags: JITFlags = NOTHING, unroll: UnrollFlags = NOTHING)[source]

Bases: ImplicitStepConfig

Configuration describing the DIRK integrator.

tableau

Butcher tableau describing the diagonally implicit method.

Type:

DIRKTableau

attempt_dense_prediction

Request dense stage prediction: accepted steps warm-start each stage’s Newton solve by reading the previous step’s stage curve ahead over the next step. Ignored when the tableau does not meet the transform’s preconditions.

Type:

bool

predictor_fn

Compiled dense-prediction device function, piped through compile settings so predictor rebuilds invalidate the step.

Type:

Callable or None

stage_increment_location

Buffer location for the working stage-increment vector.

Type:

str

stage_increment_history_location

Buffer location for the previous step’s per-stage increment history consumed by dense prediction.

Type:

str

previous_step_size_location

Buffer location for the previous-step-size scalar consumed by dense prediction.

Type:

str

stage_base_location

Buffer location for the stage base-state vector.

Type:

str

accumulator_location

Buffer location for the explicit stage accumulator.

Type:

str

stage_rhs_location

Buffer location for the cached stage effective derivative.

Type:

str

apply_mass_fn

Compiled mass-matrix product used by error smoothing.

Type:

Callable or None

inverse_mass_dxdt_fn

Compiled effective derivative M**-1 @ f for explicit stages.

Type:

Callable or None

accumulator_location: str
apply_mass_fn: Callable | None
attempt_dense_prediction: bool
inverse_mass_dxdt_fn: Callable | None
predictor_fn: Callable | None
previous_step_size_location: str
stage_base_location: str
stage_increment_history_location: str
stage_increment_location: str
stage_rhs_location: str
tableau: DIRKTableau