Simulation¶

def as_ode_term(
    sys: ContLTI | NonlinearSystem | PHSSystem,
    control_fn: Callable[[float, Array], Array],
) -> Any:
    """Wrap a continuous system as a Diffrax `ODETerm`.

    Returns an ODE term for a caller-configured Diffrax solve. The downstream
    solver and adjoint configuration determine transform behavior.

    Args:
        sys: Continuous-time system.
        control_fn: Control function mapping `(t, x)` to an input `u`.

    Returns:
        diffrax.ODETerm: An ODE term representing `x_dot = A @ x + B @ u(t, x)`.

    Raises:
        ImportError: If Diffrax is not installed.
    """
    try:
        import diffrax
    except ImportError:
        raise ImportError("diffrax is required for as_ode_term()")

    def vector_field(t, x, args):
        del args
        u = control_fn(t, x)
        if isinstance(sys, ContLTI):
            return sys.A @ x + sys.B @ u
        return sys.dynamics(t, x, u)

    return diffrax.ODETerm(vector_field)

def foh_inputs(us: Array) -> Array:
    """Build first-order-hold endpoint pairs from a sample sequence.

    Given inputs `u_k`, returns a sequence with shape `(T, 2, m)` where each
    item contains `(u_k, u_{k+1})`, using `u_{T-1}` for the final right
    endpoint.
    """
    if us.ndim < 2:
        raise ValueError("foh_inputs() expects an input array with leading time axis.")
    return jnp.stack([us, jnp.concatenate([us[1:], us[-1:]], axis=0)], axis=1)

def impulse_response(
    sys: DiscLTI | ContLTI,
    *,
    num_steps: int | None = None,
    duration: float | None = None,
    input_index: int = 0,
    x0: Array | None = None,
    dt: float | None = None,
    solver: Any = None,
    adjoint: Any = None,
    dt0: float | None = None,
    stepsize_controller: Any = None,
    max_steps: int = 100_000,
    throw: bool = False,
) -> SimResult:
    """Compute the impulse response of an LTI system.

    For discrete systems this applies a unit pulse at `k=0`. For continuous
    systems this applies the equivalent state jump `x(0+) = x0 + B[:, i]` and
    then simulates with zero input. That means the returned sampled trajectory
    omits any singular direct-feedthrough `D delta(t)` spike at `t=0`.

    The continuous-time convention represents the impulse as a state jump;
    sampled outputs omit the singular direct-feedthrough term at `t=0`.

    Args:
        sys: Continuous or discrete-time system.
        num_steps: Discrete-only number of time steps.
        duration: Continuous-only response duration.
        input_index: Input channel receiving the impulse.
        x0: Optional initial state. Defaults to zeros.
        dt: Continuous-only output sample spacing hint.
        solver: Continuous-only Diffrax solver override.
        adjoint: Continuous-only Diffrax adjoint override.
        dt0: Continuous-only initial solver step size hint.
        stepsize_controller: Native Diffrax step-size controller.
        max_steps: Continuous integration step limit.
        throw: Raise on numerical failure.

    Returns:
        SimResult: Simulation result with output times, states, and outputs.
    """
    from contrax.sim import lsim, simulate

    n = sys.A.shape[0]
    x0 = jnp.zeros(n, dtype=sys.A.dtype) if x0 is None else x0
    u_impulse = _unit_input(sys, input_index)

    if isinstance(sys, DiscLTI):
        if (
            solver is not None
            or adjoint is not None
            or dt0 is not None
            or dt is not None
            or stepsize_controller is not None
            or max_steps != 100_000
        ):
            raise ValueError("Diffrax solver configuration is continuous-only")
        if duration is not None:
            raise ValueError(
                "impulse_response() for DiscLTI requires num_steps, not duration."
            )
        if num_steps is None:
            raise ValueError("impulse_response() for DiscLTI requires num_steps.")
        if (
            not isinstance(num_steps, int)
            or isinstance(num_steps, bool)
            or num_steps < 0
        ):
            raise ValueError("num_steps must be a nonnegative static Python integer")
        T_int = num_steps
        us = jnp.zeros((T_int, sys.B.shape[1]), dtype=sys.B.dtype)
        if T_int:
            us = us.at[0].set(u_impulse)
        return lsim(sys, us, x0=x0, throw=throw)

    if num_steps is not None:
        raise ValueError(
            "impulse_response() for ContLTI requires duration, not num_steps."
        )
    if duration is None:
        raise ValueError("impulse_response() for ContLTI requires duration.")
    x0_impulse = x0 + sys.B[:, input_index]
    return simulate(
        sys,
        x0_impulse,
        lambda t, x: jnp.zeros(sys.B.shape[1], dtype=sys.B.dtype),
        duration=float(duration),
        dt=dt,
        solver=solver,
        adjoint=adjoint,
        dt0=dt0,
        stepsize_controller=stepsize_controller,
        max_steps=max_steps,
        throw=throw,
    )

def initial_response(
    sys: DiscLTI | ContLTI,
    x0: Array,
    *,
    num_steps: int | None = None,
    duration: float | None = None,
    dt: float | None = None,
    solver: Any = None,
    adjoint: Any = None,
    dt0: float | None = None,
    stepsize_controller: Any = None,
    max_steps: int = 100_000,
    throw: bool = False,
) -> SimResult:
    """Compute the zero-input response from a nonzero initial state.

    Args:
        sys: Continuous or discrete-time system.
        x0: Initial state. Shape: `(n,)`.
        num_steps: Discrete-only number of time steps.
        duration: Continuous-only response duration.
        dt: Continuous-only output sample spacing hint.
        solver: Continuous-only Diffrax solver override.
        adjoint: Continuous-only Diffrax adjoint override.
        dt0: Continuous-only initial solver step size hint.
        stepsize_controller: Native Diffrax step-size controller.
        max_steps: Continuous integration step limit.
        throw: Raise on numerical failure.

    Returns:
        SimResult: Simulation result with output times, states, and outputs.
    """
    from contrax.sim import simulate

    return simulate(
        sys,
        x0,
        lambda t, x: jnp.zeros(sys.B.shape[1], dtype=sys.B.dtype),
        num_steps=num_steps,
        duration=duration,
        dt=dt,
        solver=solver,
        adjoint=adjoint,
        dt0=dt0,
        stepsize_controller=stepsize_controller,
        max_steps=max_steps,
        throw=throw,
    )

def lsim(
    sys: DiscLTI | ContLTI,
    us: Array,  # (T, m) input sequence
    x0: Array | None = None,  # (n,) initial state; defaults to zeros
    *,
    ts: Array | None = None,
    input_interpolation: str = "zoh",
    throw: bool = False,
) -> SimResult:
    """Simulate a discrete or continuous LTI system from sampled open-loop inputs.

    Discrete indexing returns T inputs/outputs and T+1 states. Continuous
    systems require N sample times and inputs and return N sampled states and
    outputs, with x0 at ts[0]. Exact block exponentials implement ZOH or FOH
    inputs on each interval; nonuniform grids are supported.

    Uniform grids (within time-differencing roundoff) reuse one transition
    exponential. Forward/reverse and higher derivatives use the general-grid
    computation, retaining derivatives of every independent interval. This
    optimization targets primal trajectories, not differentiation throughput.

    The input sequence has a fixed shape and each interval uses the selected
    zero-order or first-order hold convention.

    Args:
        sys: Continuous or discrete LTI system.
        us: Input sequence. Shape: `(T, m)` discrete or `(N, m)` continuous.
        x0: Initial state. Shape: `(n,)`. Defaults to zeros.
        ts: Continuous-only sample times. Shape: `(N,)`, strictly increasing.
        input_interpolation: Continuous interval interpretation, zoh or foh.
        throw: Raise on numerical failure, otherwise inspect result.status.

    Returns:
        SimResult: Simulation result with `.ts`, `.xs`, `.ys`. Shapes:
            `(T,)`, `(T + 1, n)`, `(T, p)` for discrete systems and
            `(N,)`, `(N, n)`, `(N, p)` for continuous systems. `status` and
            `successful` report numerical validity.
            state_times aligns with xs, output_times (also ts) aligns with ys.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys = cx.dss(
        ...     jnp.array([[1.0, 0.1], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.1]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.1,
        ... )
        >>> result = cx.lsim(sys, jnp.zeros((20, 1)))
    """
    if not isinstance(sys, (DiscLTI, ContLTI)):
        raise TypeError("lsim requires a continuous or discrete LTI system")
    us = jnp.asarray(us)
    n = sys.A.shape[0]
    if x0 is None:
        x0 = jnp.zeros(n, dtype=sys.A.dtype)
    x0 = jnp.asarray(x0)
    if isinstance(sys, ContLTI):
        return continuous_lsim(sys, us, x0, ts, input_interpolation, throw)
    if ts is not None or input_interpolation != "zoh":
        raise ValueError(
            "Discrete lsim uses sys.dt and sample inputs; ts/FOH are continuous-only"
        )

    if us.ndim != 2 or us.shape[1] != sys.B.shape[1] or x0.shape != (n,):
        raise ValueError("us and x0 must have shapes (T, m) and (n,)")
    T = us.shape[0]
    dtype = jnp.result_type(sys.A, us, x0)
    x0, us = x0.astype(dtype), us.astype(dtype)
    ts = jnp.arange(T, dtype=sys.A.real.dtype) * sys.dt

    def step(x, u):
        x_next = sys.A @ x + sys.B @ u
        y = sys.C @ x + sys.D @ u
        return x_next, (x_next, y)

    _, (xs, ys) = jax.lax.scan(step, x0, us)
    xs_full = jnp.concatenate([x0[None], xs], axis=0)
    result = SimResult(
        ts,
        xs_full,
        ys,
        state_times=jnp.arange(T + 1, dtype=ts.dtype) * sys.dt,
        num_steps=jnp.asarray(T),
    )
    if throw:
        result = eqx.error_if(result, ~result.successful, "Discrete lsim failed")
    return result

def rollout(
    f: Callable[..., Any],
    x0: Any,
    us: Any,
    params: Any = None,
) -> Any:
    """Roll out arbitrary discrete-time dynamics over a fixed input sequence.

    Applies `f` to each input with `jax.lax.scan` and returns the states from
    `x0` through the final transition. Define output calculations separately
    when they are needed.

    Args:
        f: Transition function. If `params` is omitted, `f(x, u)` is called.
            Otherwise, `f(x, u, params)` is called.
        x0: Initial state. Shape: `(n,)` or any fixed-shape state pytree.
        us: Input sequence. Shape: `(T, m)` or any fixed-leading-axis input
            pytree accepted by `jax.lax.scan`.
        params: Optional static or array-valued parameters passed to `f`.

    Returns:
        PyTree: State trajectory with the same structure as x0. Every leaf
            includes the initial state and has shape `(T + 1, ...)`.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> def f(x, u, gain):
        ...     return gain * x + u
        >>> xs = cx.rollout(f, jnp.array([1.0]), jnp.zeros((10, 1)), 0.9)
    """

    def step(x, u):
        x_next = f(x, u) if params is None else f(x, u, params)
        return x_next, x_next

    _, xs = jax.lax.scan(step, x0, us)
    return jax.tree.map(
        lambda initial, history: jnp.concatenate(
            [jnp.asarray(initial)[None], history], axis=0
        ),
        x0,
        xs,
    )

def sample_system(
    sys: NonlinearSystem | PHSSystem,
    dt: float,
    *,
    input_interpolation: str = "zoh",
    method: str = "diffrax",
    n_substeps: int = 1,
    solver: Any = None,
    adjoint: Any = None,
    dt0: float | None = None,
    stepsize_controller: Any = None,
    max_steps: int = 100_000,
    throw: bool = False,
) -> NonlinearSystem:
    """Sample a continuous nonlinear system into a discrete transition model.

    Builds a discrete-time [NonlinearSystem][contrax.systems.NonlinearSystem]
    whose one-step dynamics integrate the continuous model over one sample
    interval using the chosen integrator.

    With `input_interpolation="zoh"`, each discrete input has shape `(m,)` and
    is held constant over the interval. With `input_interpolation="foh"`, each
    discrete input has shape `(2, m)` and is interpreted as the pair
    `(u_k, u_{k+1})`; use `foh_inputs()` to build those pairs from a sampled
    input sequence.

    **Integrator methods**

    `method="diffrax"` (default) uses adaptive Diffrax integration with Tsit5.
    For stiff dynamics, select an appropriate native implicit solver such as
    Kvaerno3. A native step-size controller sets integration tolerances. Compare
    execution costs at matched trajectory error before choosing a fixed method.

    `method="rk4"` uses a 4th-order Runge-Kutta scheme implemented as pure JAX
    arithmetic. The integration graph is fully static, so XLA compiles the entire
    step to a single fused kernel. 4th-order accurate: local error O(h⁵), global
    error O(h⁴) per sampling interval `dt / n_substeps`. As a fixed-step method,
    its accuracy depends on choosing `dt` and `n_substeps` for the system's time
    constants. **Not symplectic** — energy drift accumulates
    over long simulation horizons, which matters for conservative PHS systems.

    `method="sv"` uses Störmer-Verlet with Strang splitting for dissipation
    and inputs (B(h/2) ∘ A(h) ∘ B(h/2), where A is the symplectic flow and B
    is the dissipation/input flow). Use a
    [SeparablePHSSystem][contrax.systems.SeparablePHSSystem] with the default
    canonical `J`. Dissipative/input half-steps
    use second-order explicit midpoint, with FOH sampled at each half-step's
    own midpoint. The full split method has global order two for smooth data.
    For purely conservative systems (R=None, G=None) this reduces to standard
    Störmer-Verlet with no extra gradient evaluations. Per
    substep it evaluates `∇V` twice and `∇T` once, compared to four full
    dynamics evaluations for RK4. **Symplectic** — the conservative part
    preserves phase-space volume exactly regardless of step count; the
    dissipative/input half-steps can violate energy inequalities at large step
    sizes; resolve dissipation and forcing with `n_substeps`.
    Reference: Leimkuhler & Reich (2004), *Simulating
    Hamiltonian Dynamics*, Cambridge University Press.

    `method="euler"` and `method="midpoint"` are 1st- and 2nd-order fixed-step
    baselines. They follow the same pure-JAX, no-Diffrax path as `"rk4"`.

    `n_substeps` subdivides `dt` into `n_substeps` equal sub-intervals. Larger
    values improve accuracy at the cost of proportionally more dynamics
    evaluations. Diffrax chooses its own adaptive step count, so
    `n_substeps` applies to fixed-step methods.

    `solver`, `adjoint`, `dt0`, `stepsize_controller`, and a nondefault
    `max_steps` are Diffrax-only overrides; incompatible fixed-step
    configurations raise `ValueError`.

    Args:
        sys: Continuous source system. Must have ``dt=None``. ``"sv"`` further
            requires a [SeparablePHSSystem][contrax.systems.SeparablePHSSystem].
        dt: Sampling interval.
        input_interpolation: ``"zoh"`` (zero-order hold) or ``"foh"``
            (first-order hold, requires endpoint-pair inputs from
            `foh_inputs()`).
        method: Integration method. One of ``"diffrax"``, ``"rk4"``,
            ``"midpoint"``, ``"euler"``, or ``"sv"``. Default is ``"diffrax"``.
        n_substeps: Number of fixed sub-steps per sampling interval.
            Ignored by ``method="diffrax"``. Default is ``1``.
        solver: Diffrax solver override. Only valid with ``method="diffrax"``.
        adjoint: Diffrax adjoint override. Only valid with ``method="diffrax"``.
        dt0: Diffrax initial step-size hint. Only valid with ``method="diffrax"``.
        stepsize_controller: Native Diffrax step-size controller, defaulting to
            PIDController(rtol=1e-6, atol=1e-9). Diffrax-only.
        max_steps: Adaptive integration limit. Diffrax-only.
        throw: Raise on failed integration; otherwise step() returns diagnostics
            and dynamics() returns a nonfinite state on failure.

    Returns:
        [NonlinearSystem][contrax.systems.NonlinearSystem]: Discrete system
            with dt set and dynamics obtained from the chosen integrator.
            Its step(t, x, u) method returns SampleResult integration diagnostics.

    Examples:
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> field = cx.nonlinear_system(lambda t, x, u: -x + u)
        >>> sampled = cx.sample_system(field, 0.1, max_steps=256)
        >>> result = sampled.step(0.0, jnp.ones(1), jnp.zeros(1))
        >>> bool(result.successful)
        True
    """
    if sys.dt is not None:
        raise ValueError("sample_system() expects a continuous system with dt=None.")
    if input_interpolation not in {"zoh", "foh"}:
        raise ValueError("input_interpolation must be 'zoh' or 'foh'.")
    _valid_methods = {"diffrax"} | _FIXED_STEP_METHODS
    if method not in _valid_methods:
        raise ValueError(
            f"method must be one of {sorted(_valid_methods)!r}, got {method!r}."
        )
    if method in _FIXED_STEP_METHODS and (
        solver is not None
        or adjoint is not None
        or dt0 is not None
        or stepsize_controller is not None
        or max_steps != 100_000
    ):
        raise ValueError(
            f"solver, adjoint, and dt0 are Diffrax-only options and cannot be "
            f"used with method={method!r}."
        )
    if method == "sv":
        from contrax.phs import SeparablePHSSystem

        if not isinstance(sys, SeparablePHSSystem):
            raise ValueError(
                "method='sv' requires a SeparablePHSSystem with explicit V and T. "
                "Use separable_phs_system() to construct the system, or choose "
                "method='rk4' for a non-symplectic fixed-step alternative."
            )
        if sys.J is not None:
            raise ValueError("method='sv' requires the default canonical J (J=None)")
    if not isinstance(n_substeps, int) or n_substeps < 1:
        raise ValueError(f"n_substeps must be >= 1, got {n_substeps}.")

    dt_arr, invalid_period = sample_period(dt)
    dt_arr = eqx.error_if(
        dt_arr,
        invalid_period,
        "sample_system requires finite dt > 0",
    )
    if not isinstance(max_steps, int) or max_steps < 0:
        raise ValueError("max_steps must be a nonnegative integer")

    def output(t: Array | float, x: Array, u: Array) -> Array:
        t = jnp.asarray(t, dtype=x.real.dtype)
        if input_interpolation == "foh":
            return sys.output(t, x, u[0])
        return sys.output(t, x, u)

    if method in _FIXED_STEP_METHODS:
        fixed_transition = _fixed_step_dynamics(
            sys, dt_arr, method, n_substeps, input_interpolation
        )

        def integrate(t, x, u):
            next_x = fixed_transition(t, x, u)
            result = SampleResult(
                next_x,
                jnp.where(jnp.all(jnp.isfinite(next_x)), 0, 3),
                jnp.asarray(n_substeps),
            )
            return (
                eqx.error_if(result, ~result.successful, "Sample integration failed")
                if throw
                else result
            )
    else:
        try:
            import diffrax
        except ImportError:
            raise ImportError("diffrax is required for sample_system(method='diffrax')")

        _solver = diffrax.Tsit5() if solver is None else solver
        _adjoint = diffrax.DirectAdjoint() if adjoint is None else adjoint

        def integrate(t: Array | float, x: Array, u: Array) -> SampleResult:
            local_dt = dt_arr.astype(x.dtype)
            local_dt0 = local_dt if dt0 is None else jnp.asarray(dt0, dtype=x.dtype)

            def vf(tau, x_tau, args):
                base_t, u_step = args
                u_tau = _continuous_input_fn(u_step, tau, local_dt, input_interpolation)
                return sys.dynamics(base_t + tau, x_tau, u_tau)

            solution = diffrax.diffeqsolve(
                diffrax.ODETerm(vf),
                _solver,
                t0=jnp.asarray(0.0, dtype=x.dtype),
                t1=local_dt,
                dt0=local_dt0,
                y0=x,
                args=(jnp.asarray(t, dtype=x.dtype), u),
                saveat=diffrax.SaveAt(t1=True),
                stepsize_controller=(
                    diffrax.PIDController(rtol=1e-6, atol=1e-9)
                    if stepsize_controller is None
                    else stepsize_controller
                ),
                adjoint=_adjoint,
                max_steps=max_steps,
                throw=False,
            )
            result = SampleResult(
                solution.ys[0], _diffrax_status(solution), solution.stats["num_steps"]
            )
            return (
                eqx.error_if(result, ~result.successful, "Sample integration failed")
                if throw
                else result
            )

    def dynamics(t, x, u):
        result = integrate(t, x, u)
        return jnp.where(result.successful, result.x, jnp.full_like(result.x, jnp.nan))

    return NonlinearSystem(
        dynamics=dynamics,
        observation=output,
        dt=dt_arr,
        state_dim=sys.state_dim,
        input_dim=sys.input_dim,
        _input_shape=(2, 0 if sys.input_dim is None else sys.input_dim)
        if input_interpolation == "foh"
        else None,
        output_dim=sys.output_dim,
        _step_fn=integrate,
    )

def simulate(
    sys: DiscLTI | ContLTI | NonlinearSystem | PHSSystem,
    x0: Array,  # (n,)
    policy: Callable[[float, Array], Array],  # (t, x) -> u
    *,
    num_steps: int | None = None,
    duration: float | None = None,
    dt: float | None = None,
    solver: Any = None,
    adjoint: Any = None,
    dt0: float | None = None,
    stepsize_controller: Any = None,
    max_steps: int = 100_000,
    throw: bool = False,
) -> SimResult:
    """Simulate a system in closed loop under a control policy.

    Evaluates `policy(t, x)` along the system trajectory and dispatches to the
    simulation path for the system's time domain.

    Discrete systems use `lax.scan` and require `num_steps`. Continuous systems
    use Diffrax and require `duration`; their trajectories are sampled on a
    fixed uniform save grid.

    Continuous integration accepts native Diffrax solvers, adjoints and
    step-size controllers, plus an initial step and integration limit.
    The default continuous solver is
    `diffrax.Tsit5()` with `diffrax.RecursiveCheckpointAdjoint()`.

    Args:
        sys: Continuous or discrete LTI, nonlinear, or port-Hamiltonian system.
        x0: Initial state. Shape: `(n,)`.
        policy: Control policy mapping `(t, x)` to an input `u`.
        num_steps: Discrete-only number of time steps.
        duration: Continuous-only simulation duration.
        dt: Continuous-only output sample spacing hint. When omitted, the
            continuous path returns 201 samples including `t=0` and `t=T`.
        solver: Continuous-only Diffrax solver override.
        adjoint: Continuous-only Diffrax adjoint override.
        dt0: Continuous-only initial solver step size hint.
        stepsize_controller: Native Diffrax step-size controller. Defaults to
            PIDController(rtol=1e-6, atol=1e-9).
        max_steps: Continuous adaptive integration step limit.
        throw: Raise on numerical failure, otherwise return status diagnostics.

    Returns:
        SimResult: Simulation result. Discrete trajectories have shapes `(T,)`,
        `(T + 1, n)`, and `(T, p)`; continuous trajectories have shapes
        `(N,)`, `(N, n)`, and `(N, p)`, where `N` is the number of saved samples.
        Use state_times with xs and output_times with ys; the discrete terminal
        state has its own time. final_state returns the last state.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys_d = cx.dss(
        ...     jnp.array([[1.0, 0.05], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.05]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.05,
        ... )
        >>> K = cx.lqr(sys_d, jnp.eye(2), jnp.array([[1.0]])).K
        >>> result = cx.simulate(
        ...     sys_d,
        ...     jnp.array([1.0, 0.0]),
        ...     lambda t, x: -K @ x,
        ...     num_steps=60,
        ... )
    """
    if not isinstance(max_steps, int) or max_steps < 0:
        raise ValueError("max_steps must be a nonnegative integer")
    discrete = isinstance(sys, DiscLTI) or (
        isinstance(sys, (NonlinearSystem, PHSSystem)) and sys.dt is not None
    )
    if discrete:
        if (
            solver is not None
            or adjoint is not None
            or dt0 is not None
            or dt is not None
            or stepsize_controller is not None
            or max_steps != 100_000
        ):
            raise ValueError("Diffrax solver configuration is continuous-only")
        if num_steps is not None and (
            not isinstance(num_steps, int)
            or isinstance(num_steps, bool)
            or num_steps < 0
        ):
            raise ValueError("num_steps must be a nonnegative static Python integer")
    if isinstance(sys, DiscLTI):
        if duration is not None:
            raise ValueError("simulate() for DiscLTI requires num_steps, not duration.")
        if num_steps is None:
            raise ValueError("simulate() for DiscLTI requires num_steps.")
        result = SimResult(
            *_simulate_discrete(sys, x0, policy, num_steps),
            state_times=jnp.arange(num_steps + 1, dtype=x0.real.dtype)
            * jnp.asarray(sys.dt, dtype=x0.real.dtype),
            num_steps=jnp.asarray(num_steps),
        )
        return (
            eqx.error_if(result, ~result.successful, "Discrete simulation failed")
            if throw
            else result
        )
    if isinstance(sys, (NonlinearSystem, PHSSystem)) and sys.dt is not None:
        if duration is not None:
            raise ValueError(
                "simulate() for discrete nonlinear systems requires "
                "num_steps, not duration."
            )
        if num_steps is None:
            raise ValueError(
                "simulate() for discrete nonlinear systems requires num_steps."
            )
        result = _simulate_nonlinear_discrete(sys, x0, policy, num_steps)
        return (
            eqx.error_if(result, ~result.successful, "Discrete simulation failed")
            if throw
            else result
        )
    if num_steps is not None:
        if isinstance(sys, ContLTI):
            raise ValueError("simulate() for ContLTI requires duration, not num_steps.")
        raise ValueError(
            "simulate() for continuous nonlinear systems requires duration, "
            "not num_steps."
        )
    if duration is None:
        if isinstance(sys, ContLTI):
            raise ValueError("simulate() for ContLTI requires duration.")
        raise ValueError(
            "simulate() for continuous nonlinear systems requires duration."
        )
    return _simulate_continuous(
        sys,
        x0,
        policy,
        float(duration),
        dt=dt,
        solver=solver,
        adjoint=adjoint,
        dt0=dt0,
        stepsize_controller=stepsize_controller,
        max_steps=max_steps,
        throw=throw,
    )