physicskit 0.1.0__py3-none-any.whl

physicskit/chaos/visualizers/bifurcation.py

@@ -0,0 +1,273 @@
+"""Generic bifurcation diagrams: sweep a parameter, plot the attractor cross-section.
+
+A bifurcation diagram sweeps one parameter of a system and, for each value,
+plots the "surviving" long-term values of some observable after discarding
+an initial transient: a single point for a stable fixed point, a handful of
+points for a period-``k`` cycle, and a dense band for chaos. The core
+:func:`bifurcation_diagram` function is deliberately sampling-agnostic (it
+just calls a user-supplied ``sample_fn(param) -> ndarray`` for each parameter
+value), so it works equally well for:
+
+- **Discrete maps** (:class:`physicskit.chaos.core.base_system.DiscreteMap`), where
+  the natural "sample" is the tail of the map's own iterated trajectory --
+  see :func:`map_bifurcation_sampler`.
+- **Continuous, periodically-forced flows**
+  (:class:`physicskit.chaos.core.base_system.DynamicalSystem`), where the natural
+  "sample" is a *stroboscopic* Poincare section: the state sampled once per
+  forcing period, which turns the continuous flow into an effective discrete
+  map -- see :func:`stroboscopic_bifurcation_sampler`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, cast
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from numpy.typing import ArrayLike, NDArray
+from scipy.integrate import solve_ivp
+from tqdm.auto import tqdm
+
+from physicskit.chaos.core.base_system import DiscreteMap, DynamicalSystem
+from physicskit.chaos.visualizers import theme
+
+#: A function mapping one parameter value to the array of long-term
+#: "surviving" observable values sampled for that parameter.
+SampleFunc = Callable[[float], NDArray[np.float64]]
+
+
+def bifurcation_diagram(
+    param_values: ArrayLike,
+    sample_fn: SampleFunc,
+    n_jobs: int = 1,
+    show_progress: bool = False,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Generate raw bifurcation-diagram data by sweeping a parameter.
+
+    Parameters
+    ----------
+    param_values : array_like of float, shape (n_params,)
+        Parameter values to sweep.
+    sample_fn : callable
+        Function ``sample_fn(param) -> ndarray`` returning the array of
+        long-term ("surviving", post-transient) observable values for one
+        parameter value. See :func:`map_bifurcation_sampler` and
+        :func:`stroboscopic_bifurcation_sampler` for ready-made samplers.
+    n_jobs : int, default 1
+        Number of worker threads to evaluate `sample_fn` with, via
+        :class:`concurrent.futures.ThreadPoolExecutor`. Each parameter value
+        is fully independent, so this is embarrassingly parallel -- but the
+        speedup depends on `sample_fn` releasing the GIL while it runs:
+        :func:`stroboscopic_bifurcation_sampler` does (via
+        :func:`scipy.integrate.solve_ivp`), and so does
+        :func:`map_bifurcation_sampler` (the underlying Numba kernels are
+        compiled with ``nogil=True``); an arbitrary pure-Python `sample_fn`
+        will not see a speedup. ``1`` (the default) runs sequentially with no
+        thread-pool overhead.
+    show_progress : bool, default False
+        Display a `tqdm` progress bar over parameter values.
+
+    Returns
+    -------
+    params : ndarray of float, shape (N,)
+        Parameter value repeated once per sample.
+    samples : ndarray of float, shape (N,)
+        The concatenated samples from every parameter value, aligned with
+        `params` (``N`` is the total sample count summed across all
+        parameter values, not necessarily ``n_params`` times a fixed count).
+    """
+    params = np.atleast_1d(np.asarray(param_values, dtype=np.float64))
+
+    if n_jobs == 1:
+        results = [sample_fn(float(p)) for p in tqdm(params, disable=not show_progress, desc="sweeping")]
+    else:
+        with ThreadPoolExecutor(max_workers=n_jobs) as pool:
+            results = list(
+                tqdm(
+                    pool.map(lambda p: sample_fn(float(p)), params),
+                    total=params.shape[0],
+                    disable=not show_progress,
+                    desc="sweeping",
+                )
+            )
+
+    all_params = []
+    all_samples = []
+    for p, raw_samples in zip(params, results):
+        samples = np.atleast_1d(np.asarray(raw_samples, dtype=np.float64))
+        all_params.append(np.full(samples.shape[0], p, dtype=np.float64))
+        all_samples.append(samples)
+    return np.concatenate(all_params), np.concatenate(all_samples)
+
+
+def map_bifurcation_sampler(
+    make_map: Callable[[float], DiscreteMap],
+    state0: ArrayLike,
+    n_transient: int,
+    n_keep: int,
+    component: int = 0,
+) -> SampleFunc:
+    """Build a :func:`bifurcation_diagram` sampler for a :class:`~physicskit.chaos.core.base_system.DiscreteMap`.
+
+    Parameters
+    ----------
+    make_map : callable
+        Function ``make_map(param) -> DiscreteMap`` constructing a fresh map
+        instance for one parameter value.
+    state0 : array_like of float, shape (dim,)
+        Initial state to iterate from, for every parameter value.
+    n_transient : int
+        Number of initial iterations to discard.
+    n_keep : int
+        Number of iterations to keep (and return) after the transient.
+    component : int, default 0
+        Index of the state component to sample.
+
+    Returns
+    -------
+    callable
+        A function ``sample_fn(param) -> ndarray`` of shape ``(n_keep,)``,
+        suitable as the `sample_fn` argument to :func:`bifurcation_diagram`.
+    """
+
+    def sample_fn(param: float) -> NDArray[np.float64]:
+        system = make_map(param)
+        trajectory = system.trajectory(np.asarray(state0, dtype=np.float64), n_transient + n_keep)
+        return trajectory[n_transient:, component]
+
+    return sample_fn
+
+
+def stroboscopic_bifurcation_sampler(
+    make_system: Callable[[float], DynamicalSystem],
+    state0: ArrayLike,
+    sample_period: float,
+    n_transient_periods: int,
+    n_keep_periods: int,
+    component: int = 0,
+    dt: float = 0.01,
+    rtol: float = 1e-9,
+    atol: float = 1e-11,
+) -> SampleFunc:
+    """Build a :func:`bifurcation_diagram` sampler for a periodically-forced continuous flow.
+
+    Integrates the system and records its state stroboscopically -- once
+    every `sample_period` (typically the forcing period of a driven
+    oscillator such as :class:`physicskit.chaos.systems.continuous.Duffing`) -- which
+    turns the continuous flow into an effective discrete map for bifurcation
+    purposes. Uses :func:`scipy.integrate.solve_ivp` against the system's
+    plain-Python :meth:`~physicskit.chaos.core.base_system.DynamicalSystem.rhs`
+    (rather than the Numba-jitted low-level integrators, which require a
+    jitted right-hand side that arbitrary systems may not expose), evaluated
+    exactly at the stroboscopic sample times.
+
+    Parameters
+    ----------
+    make_system : callable
+        Function ``make_system(param) -> DynamicalSystem`` constructing a
+        fresh system instance for one parameter value.
+    state0 : array_like of float, shape (dim,)
+        Initial state to integrate from, for every parameter value.
+    sample_period : float
+        Time between stroboscopic samples (typically the forcing period).
+    n_transient_periods : int
+        Number of initial stroboscopic samples to discard.
+    n_keep_periods : int
+        Number of stroboscopic samples to keep (and return) after the
+        transient.
+    component : int, default 0
+        Index of the state component to sample.
+    dt : float, default 0.01
+        Approximate internal integration step size passed to `solve_ivp` as
+        `max_step` (the adaptive solver still lands exactly on each
+        stroboscopic sample time via `t_eval`). A parameter sweep calls this
+        once per parameter value, so loosening `dt`, `rtol`, and `atol` can
+        matter a great deal for the total sweep's running time.
+    rtol, atol : float, default 1e-9, 1e-11
+        Relative and absolute error tolerances passed to `solve_ivp`.
+
+    Returns
+    -------
+    callable
+        A function ``sample_fn(param) -> ndarray`` of shape
+        ``(n_keep_periods,)``, suitable as the `sample_fn` argument to
+        :func:`bifurcation_diagram`.
+    """
+    n_periods = n_transient_periods + n_keep_periods
+    sample_times = sample_period * np.arange(1, n_periods + 1)
+
+    def sample_fn(param: float) -> NDArray[np.float64]:
+        system = make_system(param)
+
+        def _rhs(t: float, state: NDArray[np.float64]) -> NDArray[np.float64]:
+            return system.rhs(state, t)
+
+        solution = solve_ivp(
+            _rhs,
+            (0.0, sample_times[-1]),
+            np.asarray(state0, dtype=np.float64),
+            t_eval=sample_times,
+            max_step=dt,
+            rtol=rtol,
+            atol=atol,
+        )
+        stroboscopic = solution.y[component]
+        return np.asarray(stroboscopic[n_transient_periods:], dtype=np.float64)
+
+    return sample_fn
+
+
+def plot_bifurcation_diagram(
+    param_values: ArrayLike,
+    sample_fn: SampleFunc,
+    ax: Axes | None = None,
+    n_jobs: int = 1,
+    show_progress: bool = False,
+    **scatter_kwargs: Any,
+) -> tuple[Figure, Axes]:
+    """Plot a bifurcation diagram.
+
+    Parameters
+    ----------
+    param_values : array_like of float, shape (n_params,)
+        Parameter values to sweep.
+    sample_fn : callable
+        Function ``sample_fn(param) -> ndarray`` returning the long-term
+        surviving observable values for one parameter value; see
+        :func:`map_bifurcation_sampler` / :func:`stroboscopic_bifurcation_sampler`.
+    ax : matplotlib.axes.Axes, optional
+        Axes to draw into; a new figure and axes are created if omitted.
+    n_jobs : int, default 1
+        See :func:`bifurcation_diagram`.
+    show_progress : bool, default False
+        Display a `tqdm` progress bar over parameter values.
+    **scatter_kwargs
+        Additional keyword arguments forwarded to ``ax.plot``.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : matplotlib.axes.Axes
+    """
+    params, samples = bifurcation_diagram(param_values, sample_fn, n_jobs=n_jobs, show_progress=show_progress)
+    param_array = np.asarray(param_values, dtype=np.float64)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(10, 6))
+    else:
+        fig = cast(Figure, ax.figure)
+
+    kwargs: dict[str, Any] = {
+        "marker": ",",
+        "linestyle": "none",
+        "color": theme.STRUCTURE,
+        "alpha": 0.6,
+    }
+    kwargs.update(scatter_kwargs)
+    ax.plot(params, samples, **kwargs)
+    ax.set_xlim(float(np.min(param_array)), float(np.max(param_array)))
+    return fig, ax

physicskit/chaos/visualizers/divergence.py

@@ -0,0 +1,374 @@
+"""Live trajectory-separation tracking to visually demonstrate sensitivity to
+initial conditions (a positive largest Lyapunov exponent)."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from numpy.typing import ArrayLike, NDArray
+from scipy.integrate import solve_ivp
+
+from physicskit.chaos.core.base_system import BilliardSystem, DynamicalSystem
+from physicskit.chaos.utils.metrics import lyapunov_exponent_from_divergence
+from physicskit.chaos.visualizers import theme
+
+
+def trajectory_divergence(
+    system: DynamicalSystem,
+    state0: ArrayLike | None = None,
+    delta_0: float = 1e-8,
+    t_max: float = 100.0,
+    n_points: int = 2000,
+    direction: ArrayLike | None = None,
+    seed: int | None = None,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Integrate two nearby trajectories and track their separation over time.
+
+    Parameters
+    ----------
+    system : DynamicalSystem
+        System to integrate.
+    state0 : array_like of float, shape (dim,), optional
+        Initial state of the reference trajectory; defaults to
+        ``system.initial_state()``.
+    delta_0 : float, default 1e-8
+        Initial separation between the reference and perturbed trajectories.
+    t_max : float, default 100.0
+        Total integration time.
+    n_points : int, default 2000
+        Number of evenly-spaced time points to evaluate.
+    direction : array_like of float, shape (dim,), optional
+        Direction of the initial perturbation; a random unit vector is used
+        if omitted (normalized internally regardless).
+    seed : int, optional
+        Seed for the random direction generator, for reproducibility.
+
+    Returns
+    -------
+    t : ndarray of float, shape (n_points,)
+        Time points.
+    delta_t : ndarray of float, shape (n_points,)
+        Euclidean separation between the two trajectories at each time in `t`.
+    """
+    state0 = system.initial_state() if state0 is None else np.asarray(state0, dtype=np.float64)
+    rng = np.random.default_rng(seed)
+    direction = rng.normal(size=system.dim) if direction is None else np.asarray(direction, dtype=np.float64)
+    direction = direction / float(np.linalg.norm(direction))
+    state0_perturbed = state0 + delta_0 * direction
+
+    t_eval = np.linspace(0.0, t_max, n_points)
+
+    def _rhs(t: float, state: NDArray[np.float64]) -> NDArray[np.float64]:
+        return system.rhs(state, t)
+
+    sol1 = solve_ivp(_rhs, (0.0, t_max), state0, t_eval=t_eval, rtol=1e-9, atol=1e-11)
+    sol2 = solve_ivp(_rhs, (0.0, t_max), state0_perturbed, t_eval=t_eval, rtol=1e-9, atol=1e-11)
+
+    delta_t = np.asarray(np.linalg.norm(sol1.y - sol2.y, axis=0), dtype=np.float64)
+    return sol1.t, delta_t
+
+
+def plot_lyapunov_divergence(
+    system: DynamicalSystem,
+    delta_0: float = 1e-8,
+    t_max: float = 100.0,
+    n_points: int = 2000,
+    state0: ArrayLike | None = None,
+    fit_fraction: float = 0.5,
+    ax: Axes | None = None,
+    seed: int | None = None,
+) -> tuple[Figure, Axes, float]:
+    """Plot ``ln(delta_t / delta_0)`` vs. `t` for two nearby trajectories.
+
+    A straight, positively-sloped line demonstrates exponential divergence
+    (``lambda_max > 0``); the slope is estimated by a linear fit over the
+    first `fit_fraction` of the time range, before the separation saturates
+    at the attractor's scale.
+
+    Parameters
+    ----------
+    system : DynamicalSystem
+        System to integrate.
+    delta_0 : float, default 1e-8
+        Initial separation between the reference and perturbed trajectories.
+    t_max : float, default 100.0
+        Total integration time.
+    n_points : int, default 2000
+        Number of evenly-spaced time points to evaluate.
+    state0 : array_like of float, shape (dim,), optional
+        Initial state of the reference trajectory; defaults to
+        ``system.initial_state()``.
+    fit_fraction : float, default 0.5
+        Fraction of the time range (from ``t=0``) used to fit the
+        exponential-growth slope.
+    ax : matplotlib.axes.Axes, optional
+        Axes to draw into; a new figure and axes are created if omitted.
+    seed : int, optional
+        Seed for the random perturbation direction, for reproducibility.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : matplotlib.axes.Axes
+    lambda_max : float
+        Estimated largest Lyapunov exponent from the fitted slope.
+    """
+    t, delta_t = trajectory_divergence(system, state0=state0, delta_0=delta_0, t_max=t_max, n_points=n_points, seed=seed)
+    log_ratio = np.log(np.clip(delta_t, 1e-300, None) / delta_0)
+
+    n_fit = max(2, int(fit_fraction * len(t)))
+    lam = lyapunov_exponent_from_divergence(t[:n_fit], delta_t[:n_fit] / delta_0)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 5))
+    else:
+        fig = cast(Figure, ax.figure)
+
+    ax.plot(t, log_ratio, lw=1.2, color=theme.PRIMARY, label=r"$\ln(\delta_t/\delta_0)$")
+    ax.plot(
+        t[:n_fit],
+        lam * t[:n_fit],
+        "--",
+        color=theme.MUTED,
+        label=rf"fit: $\lambda_{{max}} \approx {lam:.3f}$",
+    )
+    ax.set_xlabel("t")
+    ax.set_ylabel(r"$\ln(\delta_t / \delta_0)$")
+    ax.set_title(f"{system.__class__.__name__} trajectory divergence")
+    ax.legend()
+    return fig, ax, lam
+
+
+def trajectory_ensemble(
+    system: DynamicalSystem,
+    state0: ArrayLike | None = None,
+    n_members: int = 40,
+    spread: float = 1e-4,
+    dt: float = 0.01,
+    n_steps: int = 10000,
+    seed: int | None = None,
+) -> NDArray[np.float64]:
+    """Integrate an ensemble of trajectories launched from a small ball
+    around a shared initial state.
+
+    Visualizing the result (e.g. with :func:`plot_trajectory_ensemble`)
+    makes sensitivity to initial conditions directly visible: a regular
+    orbit's members stay bundled together for the whole run, while a
+    chaotic orbit's members visibly fan out, in contrast to
+    :func:`trajectory_divergence`, which only tracks a *single* pair's
+    separation as a number.
+
+    Parameters
+    ----------
+    system : DynamicalSystem
+        System to integrate.
+    state0 : array_like of float, shape (dim,), optional
+        Center of the initial ball; defaults to ``system.initial_state()``.
+    n_members : int, default 40
+        Number of trajectories in the ensemble.
+    spread : float, default 1e-4
+        Radius of the ball each member's initial state is drawn from.
+    dt : float, default 0.01
+        Integration step size.
+    n_steps : int, default 10000
+        Number of integration steps per trajectory.
+    seed : int, optional
+        Seed for the random initial-state offsets, for reproducibility.
+
+    Returns
+    -------
+    ndarray of float, shape (n_members, n_steps + 1, dim)
+        Every member's full trajectory.
+    """
+    state0 = system.initial_state() if state0 is None else np.asarray(state0, dtype=np.float64)
+    rng = np.random.default_rng(seed)
+
+    all_states = np.empty((n_members, n_steps + 1, system.dim))
+    for i in range(n_members):
+        offset = rng.normal(size=system.dim)
+        offset *= spread / float(np.linalg.norm(offset))
+        _, states = system.trajectory(state0=state0 + offset, dt=dt, n_steps=n_steps)
+        all_states[i] = states
+    return all_states
+
+
+def plot_trajectory_ensemble(
+    system: DynamicalSystem,
+    state0: ArrayLike | None = None,
+    n_members: int = 40,
+    spread: float = 1e-4,
+    dt: float = 0.01,
+    n_steps: int = 10000,
+    x_index: int = 0,
+    y_index: int = 1,
+    seed: int | None = None,
+    ax: Axes | None = None,
+    **line_kwargs: Any,
+) -> tuple[Figure, Axes]:
+    """Plot an ensemble of nearby trajectories overlaid on the same axes.
+
+    Parameters
+    ----------
+    system : DynamicalSystem
+        System to integrate.
+    state0 : array_like of float, shape (dim,), optional
+        Center of the initial ball; defaults to ``system.initial_state()``.
+    n_members : int, default 40
+        Number of trajectories in the ensemble.
+    spread : float, default 1e-4
+        Radius of the ball each member's initial state is drawn from.
+    dt : float, default 0.01
+        Integration step size.
+    n_steps : int, default 10000
+        Number of integration steps per trajectory.
+    x_index, y_index : int, default 0, 1
+        State indices to use as the plotted x/y coordinates.
+    seed : int, optional
+        Seed for the random initial-state offsets, for reproducibility.
+    ax : matplotlib.axes.Axes, optional
+        Axes to draw into; a new figure and axes are created if omitted.
+    **line_kwargs
+        Additional keyword arguments forwarded to every ``ax.plot`` call.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : matplotlib.axes.Axes
+    """
+    all_states = trajectory_ensemble(system, state0=state0, n_members=n_members, spread=spread, dt=dt, n_steps=n_steps, seed=seed)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 6))
+    else:
+        fig = cast(Figure, ax.figure)
+
+    kwargs: dict[str, Any] = {"color": theme.PRIMARY, "lw": 0.5, "alpha": 0.5}
+    kwargs.update(line_kwargs)
+    for states in all_states:
+        ax.plot(states[:, x_index], states[:, y_index], **kwargs)
+    ax.set_title(f"{system.__class__.__name__} trajectory ensemble ({n_members} members, spread={spread:g})")
+    return fig, ax
+
+
+def billiard_trajectory_divergence(
+    billiard: BilliardSystem,
+    pos: ArrayLike,
+    vel: ArrayLike,
+    delta_0: float = 1e-8,
+    n_bounces: int = 200,
+) -> tuple[NDArray[np.int64], NDArray[np.float64]]:
+    """Trace two nearby billiard rays and track their bounce-by-bounce separation.
+
+    Unlike :func:`trajectory_divergence` (continuous flows, compared at
+    matched *times*), two billiard trajectories are naturally compared at
+    matched *bounce number*: the reference ray and a ray launched at a
+    slightly perturbed angle both hit the boundary `n_bounces` times, and
+    the Euclidean distance between their `k`-th hit points is tracked as `k`
+    increases.
+
+    Parameters
+    ----------
+    billiard : BilliardSystem
+        Billiard to trace both rays within.
+    pos : array_like of float, shape (2,)
+        Shared initial position ``(x, y)``; must lie in the billiard's
+        interior.
+    vel : array_like of float, shape (2,)
+        Initial velocity direction ``(vx, vy)`` of the reference ray;
+        normalized internally.
+    delta_0 : float, default 1e-8
+        Initial angular perturbation (radians) applied to `vel` to launch
+        the second ray.
+    n_bounces : int, default 200
+        Number of reflections to trace for each ray.
+
+    Returns
+    -------
+    bounce : ndarray of int, shape (n_bounces,)
+        Bounce index, ``1, 2, ..., n_bounces``.
+    separation : ndarray of float, shape (n_bounces,)
+        Euclidean distance between the two rays' hit points at each bounce.
+    """
+    vel_ref = np.asarray(vel, dtype=np.float64)
+    vel_ref = vel_ref / np.linalg.norm(vel_ref)
+    angle = float(np.arctan2(vel_ref[1], vel_ref[0])) + delta_0
+    vel_perturbed = np.array([np.cos(angle), np.sin(angle)])
+
+    result = billiard.simulate(pos, vel_ref, n_bounces)
+    result_perturbed = billiard.simulate(pos, vel_perturbed, n_bounces)
+
+    separation = np.hypot(result["x"] - result_perturbed["x"], result["y"] - result_perturbed["y"])
+    bounce = np.arange(1, n_bounces + 1)
+    return bounce, separation
+
+
+def plot_billiard_divergence(
+    billiard: BilliardSystem,
+    pos: ArrayLike,
+    vel: ArrayLike,
+    delta_0: float = 1e-8,
+    n_bounces: int = 200,
+    fit_fraction: float = 0.5,
+    ax: Axes | None = None,
+) -> tuple[Figure, Axes, float]:
+    """Plot ``ln(separation / delta_0)`` vs. bounce number for two nearby billiard rays.
+
+    A straight, positively-sloped line demonstrates exponential divergence
+    per bounce; the slope is estimated by a linear fit over the first
+    `fit_fraction` of the bounces, before the separation saturates at the
+    table's scale.
+
+    Parameters
+    ----------
+    billiard : BilliardSystem
+        Billiard to trace both rays within.
+    pos : array_like of float, shape (2,)
+        Shared initial position ``(x, y)``; must lie in the billiard's
+        interior.
+    vel : array_like of float, shape (2,)
+        Initial velocity direction ``(vx, vy)`` of the reference ray.
+    delta_0 : float, default 1e-8
+        Initial angular perturbation (radians) applied to `vel`.
+    n_bounces : int, default 200
+        Number of reflections to trace for each ray.
+    fit_fraction : float, default 0.5
+        Fraction of the bounces (from the first) used to fit the
+        exponential-growth slope.
+    ax : matplotlib.axes.Axes, optional
+        Axes to draw into; a new figure and axes are created if omitted.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : matplotlib.axes.Axes
+    lambda_per_bounce : float
+        Estimated divergence rate, in nats per bounce.
+    """
+    bounce, separation = billiard_trajectory_divergence(billiard, pos, vel, delta_0=delta_0, n_bounces=n_bounces)
+    log_ratio = np.log(np.clip(separation, 1e-300, None) / delta_0)
+
+    n_fit = max(2, int(fit_fraction * len(bounce)))
+    lam = lyapunov_exponent_from_divergence(bounce[:n_fit].astype(np.float64), separation[:n_fit] / delta_0)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 5))
+    else:
+        fig = cast(Figure, ax.figure)
+
+    ax.plot(bounce, log_ratio, lw=1.2, color=theme.PRIMARY, label=r"$\ln(\delta_k/\delta_0)$")
+    ax.plot(
+        bounce[:n_fit],
+        lam * bounce[:n_fit],
+        "--",
+        color=theme.MUTED,
+        label=rf"fit: $\lambda \approx {lam:.3f}$ / bounce",
+    )
+    ax.set_xlabel("bounce number k")
+    ax.set_ylabel(r"$\ln(\delta_k / \delta_0)$")
+    ax.set_title(f"{billiard.__class__.__name__} trajectory divergence")
+    ax.legend()
+    return fig, ax, lam