step_solver_jax

step_solver_jax(propagator_fn, generator, t_span, t_eval, y0, max_dt, time_grid_jitter=None, *, rng=None)

Integrate the time evolution using JAX primitives and scans.

Parameters:

  • propagator_fn (Callable) – Callable returning the evolution operator for a time step when supplied with a generator.

  • generator (Callable) – Function of time that produces the system generator matrix.

  • t_span (ndarray | Array) – Two-element array containing the start and end times.

  • t_eval (ndarray | Array | None) – Optional evaluation grid contained within t_span.

  • y0 (ndarray | Array) – Initial state or operator to propagate.

  • max_dt (float) – Maximum allowed time-step length.

  • time_grid_jitter (float | None) – Optional jitter amplitude forwarded to compute_time_grid().

Return type:

SolverResult

Returns:

Solver result storing the evaluation grid and propagated states.