Skip to content

Integrators

Numerical ODE integrators for orbital dynamics propagation.

Provides fixed-step and adaptive Runge-Kutta integrators, all implemented in JAX for compatibility with jax.jit, jax.vmap, and automatic differentiation.

Available integrators:

  • :func:rk4_step -- Classic 4th-order Runge-Kutta (fixed step)
  • :func:rkf45_step -- Runge-Kutta-Fehlberg 4(5) (adaptive step)
  • :func:dp54_step -- Dormand-Prince 5(4) (adaptive step)
  • :func:rkn1210_step -- Runge-Kutta-Nyström 12(10) (adaptive step, second-order ODE)

All step functions share a common interface::

result = step_fn(dynamics, t, state, dt)

where dynamics(t, x) -> dx defines the ODE right-hand side, and the result is a :class:StepResult named tuple.

AdaptiveConfig

Bases: NamedTuple

Configuration for adaptive step-size control.

Used by rkf45_step and dp54_step to control step acceptance, rejection, and step-size adjustment. Default values provide a reasonable starting point for orbital mechanics problems.

Attributes:

Name Type Description
abs_tol float

Absolute error tolerance per component. Components with magnitude near zero are controlled by this tolerance.
rel_tol float

Relative error tolerance per component. Components with large magnitude are controlled by this tolerance.
safety_factor float

Multiplicative safety factor applied to step-size predictions. Values < 1.0 produce conservative step sizes.
min_scale_factor float

Minimum allowed ratio dt_next / dt_used. Prevents excessively aggressive step-size reduction.
max_scale_factor float

Maximum allowed ratio dt_next / dt_used. Prevents excessively aggressive step-size growth.
min_step float

Absolute minimum allowed step size. If the adaptive algorithm would reduce below this, the step is accepted regardless of error.
max_step float

Absolute maximum allowed step size.
max_step_attempts int

Maximum number of step-rejection retries before accepting the step regardless. Prevents infinite loops.

StepResult

Bases: NamedTuple

Result of a single integrator step.

Returned by all step functions (rk4_step, rkf45_step, dp54_step). For fixed-step methods (RK4), error_estimate is always 0.0 and dt_next equals dt_used.

Attributes:

Name Type Description
state Array

State vector at time t + dt_used.
dt_used Array

Actual timestep taken. For adaptive methods, this may be smaller than the requested dt if the step was rejected and retried.
error_estimate Array

Normalized error estimate. A value <= 1.0 means the step met the tolerance. Always 0.0 for RK4.
dt_next Array

Suggested timestep for the next step. For adaptive methods, this is computed from the error estimate. For RK4, equals dt_used.

dp54_step(dynamics, t, state, dt, config=None, control=None)

Perform a single adaptive DP54 integration step.

Advances the state from time t by up to dt using the Dormand-Prince 5(4) method with adaptive step-size control. If the error exceeds the tolerance, the step is rejected and retried with a smaller timestep.

Compatible with jax.jit and jax.vmap. Not compatible with reverse-mode jax.grad due to the internal lax.while_loop.

Parameters:

Name Type Description Default
dynamics Callable[[ArrayLike, ArrayLike], Array]

ODE right-hand side function f(t, x) -> dx/dt.

 required
t ArrayLike

Current time.

 required
state ArrayLike

Current state vector.

 required
dt ArrayLike

Requested timestep. May be negative for backward integration. The actual timestep used may be smaller if the adaptive controller rejects the initial attempt.

 required
config AdaptiveConfig | None

Adaptive step-size configuration. Uses default :class:AdaptiveConfig if None.

 None
control Callable[[ArrayLike, ArrayLike], Array] | None

Optional additive control function u(t, x) -> force. When provided, the effective derivative is f(t, x) + u(t, x).

 None

Returns:

Name Type Description
StepResult StepResult

Named tuple with fields: - state: State at t + dt_used. - dt_used: Actual timestep taken (<= |dt|). - error_estimate: Normalized error of the accepted step. - dt_next: Suggested timestep for the next step.

Examples:

import jax.numpy as jnp
from astrojax.integrators import dp54_step
def harmonic(t, x):
    return jnp.array([x[1], -x[0]])
result = dp54_step(harmonic, 0.0, jnp.array([1.0, 0.0]), 0.1)
result.state  # ~[cos(0.1), -sin(0.1)]

rk4_step(dynamics, t, state, dt, control=None)

Perform a single RK4 integration step.

Advances the state from time t to t + dt using the classic 4th-order Runge-Kutta method. Compatible with jax.jit and jax.vmap.

Parameters:

Name Type Description Default
dynamics Callable[[ArrayLike, ArrayLike], Array]

ODE right-hand side function f(t, x) -> dx/dt.

 required
t ArrayLike

Current time.

 required
state ArrayLike

Current state vector.

 required
dt ArrayLike

Timestep to take. May be negative for backward integration.

 required
control Callable[[ArrayLike, ArrayLike], Array] | None

Optional additive control function u(t, x) -> force. When provided, the effective derivative is f(t, x) + u(t, x).

 None

Returns:

Name Type Description
StepResult StepResult

Named tuple with fields: - state: State at t + dt. - dt_used: Always equals dt. - error_estimate: Always 0.0 (no error estimate for fixed-step methods). - dt_next: Always equals dt.

Examples:

import jax.numpy as jnp
from astrojax.integrators import rk4_step
def harmonic(t, x):
    return jnp.array([x[1], -x[0]])
result = rk4_step(harmonic, 0.0, jnp.array([1.0, 0.0]), 0.01)
result.state  # ~[cos(0.01), -sin(0.01)]

rkf45_step(dynamics, t, state, dt, config=None, control=None)

Perform a single adaptive RKF45 integration step.

Advances the state from time t by up to dt using the Runge-Kutta-Fehlberg 4(5) method with adaptive step-size control. If the error exceeds the tolerance, the step is rejected and retried with a smaller timestep.

Compatible with jax.jit and jax.vmap. Not compatible with reverse-mode jax.grad due to the internal lax.while_loop.

Parameters:

Name Type Description Default
dynamics Callable[[ArrayLike, ArrayLike], Array]

ODE right-hand side function f(t, x) -> dx/dt.

 required
t ArrayLike

Current time.

 required
state ArrayLike

Current state vector.

 required
dt ArrayLike

Requested timestep. May be negative for backward integration. The actual timestep used may be smaller if the adaptive controller rejects the initial attempt.

 required
config AdaptiveConfig | None

Adaptive step-size configuration. Uses default :class:AdaptiveConfig if None.

 None
control Callable[[ArrayLike, ArrayLike], Array] | None

Optional additive control function u(t, x) -> force. When provided, the effective derivative is f(t, x) + u(t, x).

 None

Returns:

Name Type Description
StepResult StepResult

Named tuple with fields: - state: State at t + dt_used. - dt_used: Actual timestep taken (<= |dt|). - error_estimate: Normalized error of the accepted step. - dt_next: Suggested timestep for the next step.

Examples:

import jax.numpy as jnp
from astrojax.integrators import rkf45_step
def harmonic(t, x):
    return jnp.array([x[1], -x[0]])
result = rkf45_step(harmonic, 0.0, jnp.array([1.0, 0.0]), 0.1)
result.state  # ~[cos(0.1), -sin(0.1)]

rkn1210_step(dynamics, t, state, dt, config=None, control=None)

Perform a single adaptive RKN1210 integration step.

Advances the state from time t by up to dt using the Runge-Kutta-Nyström 12(10) method with adaptive step-size control. This integrator exploits the second-order structure of y'' = f(t, y) problems internally, while accepting the standard dynamics(t, state) interface.

The state vector must have an even number of elements, split as [position, velocity]. The dynamics function must return [velocity, acceleration]. Only the acceleration half is used internally for stage computation.

Compatible with jax.jit and jax.vmap. Not compatible with reverse-mode jax.grad due to the internal lax.while_loop.

Parameters:

Name Type Description Default
dynamics Callable[[ArrayLike, ArrayLike], Array]

ODE right-hand side function f(t, x) -> dx/dt. Must return [velocity, acceleration] for a state vector [position, velocity].

 required
t ArrayLike

Current time.

 required
state ArrayLike

Current state vector. Must be even-length with the first half representing positions and the second half velocities.

 required
dt ArrayLike

Requested timestep. May be negative for backward integration. The actual timestep used may be smaller if the adaptive controller rejects the initial attempt.

 required
config AdaptiveConfig | None

Adaptive step-size configuration. Uses default :class:AdaptiveConfig if None.

 None
control Callable[[ArrayLike, ArrayLike], Array] | None

Optional additive control function u(t, x) -> force. When provided, the effective derivative is f(t, x) + u(t, x).

 None

Returns:

Name Type Description
StepResult StepResult

Named tuple with fields: - state: State at t + dt_used. - dt_used: Actual timestep taken (<= |dt|). - error_estimate: Normalized error of the accepted step. - dt_next: Suggested timestep for the next step.

Examples:

import jax.numpy as jnp
from astrojax.integrators import rkn1210_step
def harmonic(t, x):
    return jnp.array([x[1], -x[0]])
result = rkn1210_step(harmonic, 0.0, jnp.array([1.0, 0.0]), 0.1)
result.state  # ~[cos(0.1), -sin(0.1)]