ott.tools.plot.transport_animation

ott.tools.plot.transport_animation#

ott.tools.plot.transport_animation(n_frames, static_src_points, static_tgt_points, *, n_grid=0, velocity_field=None, dynamic_src_points=None, num_ifm_interpolants=0, plot_ifm_arrows=False, title=None, figsize=(8, 6), xlimits=None, ylimits=None, padding=0.1, interval=300, save_path=None, darkmode=False)[source]#

Create animated visualizations of optimal transport and flow matching.

This function generates animations illustrating various aspects of optimal transport, including :term`Monge map`, McCann interpolation, Benamou-Brenier velocity fields, and flow matching approaches. It supports multiple visualization modes and can display static point clouds, dynamic trajectories, and velocity fields on grids.

Parameters:

  • n_frames (int) – Number of animation frames \(> 1\). If 1, creates a static plot instead of an animation.

  • static_src_points (Array) – Source distribution points of shape [n, 2], representing samples from \(\mu_0\).

  • static_tgt_points (Array) – Target distribution points of shape [n, 2], representing samples from \(\mu_1\).

  • n_grid (int) – Number of grid points per dimension for displaying velocity fields. If \(> 0\), displays velocity field arrows on a uniform \(n_{grid} \times n_{grid}\) grid.

  • velocity_field (Optional[Callable[[Array, Array], Array]]) – Optional learned/estimated velocity field function with a signature v(t, x) -> velocity where t is time of shape [batch,] and x is position of shape [batch, 2].

  • dynamic_src_points (Optional[Array]) – Additional points of shape [m, 2] to highlight and transport dynamically through the animation. Useful for emphasizing specific trajectories. If None and the velocity field is specified, use static_src_points.

  • num_ifm_interpolants (Union[int, Literal['all']]) – Number of random pairs of source and target points for the independent flow matching (IFM) interpolant \((1-t)x_0 + tx_1\). If 'all', use len(static_src_points) ** 2.

  • plot_ifm_arrows (bool) – Whether to display velocity arrows for the IFM interpolant. Only relevant when num_ifm_interpolants > 0.

  • title (Optional[str]) – Title for the animation/plot.

  • figsize (Tuple[int, int]) – Figure size as (width, height) in inches.

  • xlimits (Optional[Tuple[float, float]]) – X-axis limits as (xmin, xmax). If None, it is computed automatically from all points with padding.

  • ylimits (Optional[Tuple[float, float]]) – Y-axis limits as (ymin, ymax). If None, it is computed automatically from all points with padding.

  • padding (float) – Fractional padding to add around automatically computed axis limits. For example, 0.1 adds 10% padding on each side.

  • interval (int) – Used for animations, delay between frames in milliseconds.

  • save_path (Optional[str]) – Path where to save the animation/plot.

  • darkmode (bool) – whether plotting is done in darkmode.

Return type:

Union[Figure, FuncAnimation]

Returns:

An animation object if n_frames > 1, otherwise a figure.