FIRKStep

FIRKStep provides a fully implicit Runge–Kutta integrator that solves coupled stage systems with the cached Newton–Krylov helpers supplied by cubie.integrators.matrix_free_solvers. The factory consumes FIRKTableau instances, exposing high-order Gauss–Legendre and Radau IIA schemes for stiff problems while preserving adaptive error control through embedded weights.

Defaults

algorithm="firk" integrates with the firk_gauss_legendre_2 tableau (two-stage Gauss–Legendre, order 4), solving all stages together as one coupled Newton–Krylov system with the shared defaults listed in Default settings. The default tableau has no embedded error estimate, so it runs under fixed-step control; tableaus that provide one — radau (Radau IIA-5), for example — enable the family’s Gustafsson predictive defaults (step-size growth clamped to 0.2–8.0×). Named schemes are listed on the FIRK tableau registry page.

class cubie.integrators.algorithms.FIRKStep(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: FIRKTableau = FIRKTableau(a=((0.25, np.float64(-0.038675134594812866)), (np.float64(0.5386751345948129), 0.25)), b=(0.5, 0.5), c=(np.float64(0.21132486540518713), np.float64(0.7886751345948129)), order=4, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=8.0, dense_prediction_ratio_float64=8.0, defaults={}), n_drivers: int = 0, attempt_dense_prediction: bool = True, **kwargs)

Bases: ODEImplicitStep

Fully 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: FIRKTableau = FIRKTableau(a=((0.25, np.float64(-0.038675134594812866)), (np.float64(0.5386751345948129), 0.25)), b=(0.5, 0.5), c=(np.float64(0.21132486540518713), np.float64(0.7886751345948129)), order=4, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=8.0, dense_prediction_ratio_float64=8.0, defaults={}), n_drivers: int = 0, attempt_dense_prediction: bool = True, **kwargs) → None

Initialise the FIRK step configuration.

This constructor creates a FIRK 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 (Gustafsson controller), 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 – FIRK tableau describing the coefficients. Defaults to DEFAULT_FIRK_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 FIRKStepConfig, 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 FIRK_ADAPTIVE_DEFAULTS (Gustafsson controller)

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

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

FIRK methods require solving a coupled system of all stages simultaneously, which is more computationally expensive than DIRK methods but can achieve higher orders of accuracy for stiff systems.

use_smoothed_error defaults on when the tableau supports it.

_build_error_solver() → None

Construct the width-n smoothing solver from live settings.

property algorithm_order: int

Return the order of accuracy used for step-size control.

property baked_stage_diagonal: None

smoothing solves use transform eigenvalues.

Type:

Return None

build_implicit_helpers() → None

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

Compile the FIRK device step.

default_tableau = FIRKTableau(a=((0.25, np.float64(-0.038675134594812866)), (np.float64(0.5386751345948129), 0.25)), b=(0.5, 0.5), c=(np.float64(0.21132486540518713), np.float64(0.7886751345948129)), order=4, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=8.0, dense_prediction_ratio_float64=8.0, defaults={})

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

classmethod family_defaults(tableau=None) → AlgorithmDefaults

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 optimisation_candidates: Tuple[Dict[str, Any], ...]

Newton unrolling crossed with stage_increment placement, plus rolled other_small at rolled Newton per placement.

property order: int

Return the classical order of accuracy.

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

Share a Krylov solve’s stage_increment while the GPU stays full.

register_buffers() → None

Register buffers according to locations in compile settings.

property stage_count: int

Return the number of stages described by the tableau.

property threads_per_step: int

Return the number of CUDA threads that advance one state.

class cubie.integrators.algorithms.generic_firk.FIRKStepConfig(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: FIRKTableau = FIRKTableau(a=((0.25, np.float64(-0.038675134594812866)), (np.float64(0.5386751345948129), 0.25)), b=(0.5, 0.5), c=(np.float64(0.21132486540518713), np.float64(0.7886751345948129)), order=4, b_hat=None, embedded_order=None, dense_prediction_ratio_float16=0.0, dense_prediction_ratio_float32=8.0, dense_prediction_ratio_float64=8.0, defaults={}), attempt_dense_prediction: bool = True, predictor_fn: Callable | None = None, stage_increment_location: str = 'local', previous_step_size_location: str = 'local', stage_driver_stack_location: str = 'local', stage_state_location: str = 'local', apply_mass_fn: Callable | None = None, *, jit_flags: JITFlags = NOTHING, unroll: UnrollFlags = NOTHING)

Bases: ImplicitStepConfig

Configuration describing the FIRK integrator.

tableau

Butcher tableau describing the fully implicit method.

Type:

FIRKTableau

attempt_dense_prediction

Request dense stage prediction: accepted steps warm-start Newton 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 coupled stage-increment vector.

Type:

str

previous_step_size_location

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

Type:

str

stage_driver_stack_location

Buffer location for the per-stage driver samples.

Type:

str

stage_state_location

Buffer location for the stage-state scratch vector.

Type:

str

apply_mass_fn: Callable | None

attempt_dense_prediction: bool

predictor_fn: Callable | None

previous_step_size_location: str

property smoothed_error_weights: tuple

Return the smoothed error weights cast to precision.

property solver_width: int

Return the coupled solver width across all stages.

property stage_count: int

Return the number of stages described by the tableau.

stage_driver_stack_location: str

stage_increment_location: str

stage_state_location: str

tableau: FIRKTableau