Animation Module

animation

Animation utilities for visualizing lattice field evolution.

This package provides tools to create animations of scalar field dynamics across different dimensionalities (1D, 2D, 3D), phase space visualizations, and comparison plots.

Functions are organized by category:

• 1d: 1D field animations and phase space
• 2d: 2D field animations and phase portraits
• 3d: 3D slice and isosurface animations
• comparison: Side-by-side and dimensional comparisons

animate_multi_snapshot_grid(times, snapshots, potential=None, n_snapshots=6, interval=100, save_path=None, figsize=(16, 10), title='Field Evolution Snapshots', cmap='RdBu_r', show_energy=True)

Create an animation showing a grid of snapshots at different times.

This provides a "time-lapse" view where you can see multiple moments of the evolution simultaneously. Particularly useful for understanding the overall dynamics and identifying key transition moments.

Parameters:

Name	Type	Description	Default
times	Array	Array of snapshot times	required
snapshots	list[Lattice]	List of Lattice objects	required
potential	Callable | None	Potential function for energy computation	None
n_snapshots	int	Number of time snapshots to show in grid (2, 4, 6, or 9)	6
interval	int	Delay between frames in milliseconds	100
save_path	str | None	If provided, save animation to this path	None
figsize	tuple[int, int]	Figure size	(16, 10)
title	str	Animation title	'Field Evolution Snapshots'
cmap	str	Colormap for 2D/3D visualizations	'RdBu_r'
show_energy	bool	Whether to include energy evolution plot	True

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Example

# Show 6 snapshots evolving together
anim = animate_multi_snapshot_grid(
    times, snapshots, potential=pot,
    n_snapshots=6, save_path="timelapse.gif"
)

Source code in jaxlatt/animation/comparison.py

def animate_multi_snapshot_grid(
    times: Array,
    snapshots: list[Lattice],
    potential: Callable | None = None,
    n_snapshots: int = 6,
    interval: int = 100,
    save_path: str | None = None,
    figsize: tuple[int, int] = (16, 10),
    title: str = "Field Evolution Snapshots",
    cmap: str = "RdBu_r",
    show_energy: bool = True,
) -> FuncAnimation:
    """
    Create an animation showing a grid of snapshots at different times.

    This provides a "time-lapse" view where you can see multiple moments
    of the evolution simultaneously. Particularly useful for understanding
    the overall dynamics and identifying key transition moments.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects
        potential: Potential function for energy computation
        n_snapshots: Number of time snapshots to show in grid (2, 4, 6, or 9)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title
        cmap: Colormap for 2D/3D visualizations
        show_energy: Whether to include energy evolution plot

    Returns:
        FuncAnimation object

    Example:
        ```python
        # Show 6 snapshots evolving together
        anim = animate_multi_snapshot_grid(
            times, snapshots, potential=pot,
            n_snapshots=6, save_path="timelapse.gif"
        )
        ```
    """
    lattice = snapshots[0]
    ndim = lattice.ndim

    # Determine grid layout
    if n_snapshots == 2:
        nrows, ncols = 1, 2
    elif n_snapshots == 4:
        nrows, ncols = 2, 2
    elif n_snapshots == 6:
        nrows, ncols = 2, 3
    elif n_snapshots == 9:
        nrows, ncols = 3, 3
    else:
        raise ValueError("n_snapshots must be 2, 4, 6, or 9")

    # Create figure
    if show_energy and potential is not None:
        fig = plt.figure(figsize=figsize)
        gs = GridSpec(nrows + 1, ncols, figure=fig, height_ratios=[1] * nrows + [0.4])
        axes_grid = [[fig.add_subplot(gs[i, j]) for j in range(ncols)] for i in range(nrows)]
        ax_energy = fig.add_subplot(gs[nrows, :])

        # Compute energy
        energy_data = [energy_components_integrated(s, potential) for s in snapshots]
        E_total = [e["total"] for e in energy_data]
    else:
        fig, axes_grid = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows == 1:
            axes_grid = [axes_grid]
        ax_energy = None

    axes = [ax for row in axes_grid for ax in row]

    # Initialize plots based on dimensionality
    plot_objects = []
    time_texts = []

    if ndim == 1:
        # 1D: Line plots
        length = lattice.length[0] if isinstance(lattice.length, tuple) else lattice.length
        size = lattice.size[0] if isinstance(lattice.size, tuple) else lattice.size
        x = np.linspace(0, float(length), int(size), endpoint=False)

        for ax in axes:
            (line,) = ax.plot([], [], "b-", linewidth=2)
            plot_objects.append(line)
            ax.set_xlim(0, float(length))
            ax.grid(True, alpha=0.3)

            time_text = ax.text(
                0.02,
                0.95,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
            )
            time_texts.append(time_text)

    elif ndim == 2:
        # 2D: Heatmaps
        extent = [0, lattice.length[0], 0, lattice.length[1]]

        for ax in axes:
            im = ax.imshow(
                np.asarray(snapshots[0].field).T,
                origin="lower",
                cmap=cmap,
                extent=extent,
                animated=True,
            )
            plot_objects.append(im)
            ax.set_aspect("equal")

            time_text = ax.text(
                0.02,
                0.98,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                color="white",
                fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
            )
            time_texts.append(time_text)

    elif ndim == 3:
        # 3D: Central slices
        mid_z = lattice.size[2] // 2
        extent = [0, lattice.length[0], 0, lattice.length[1]]

        for ax in axes:
            im = ax.imshow(
                np.asarray(snapshots[0].field)[:, :, mid_z].T,
                origin="lower",
                cmap=cmap,
                extent=extent,
                animated=True,
            )
            plot_objects.append(im)
            ax.set_aspect("equal")
            ax.set_title("XY slice (z=mid)", fontsize=9)

            time_text = ax.text(
                0.02,
                0.98,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                color="white",
                fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
            )
            time_texts.append(time_text)

    # Energy plot initialization
    energy_marker = None
    if ax_energy is not None:
        (line_E,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total Energy")
        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)
        ax_energy.set_xlim(times[0], times[-1])
        ax_energy.set_ylim(min(E_total) * 0.9, max(E_total) * 1.1)
        ax_energy.set_xlabel("Time", fontsize=11)
        ax_energy.set_ylabel("Energy", fontsize=11)
        ax_energy.legend()
        ax_energy.grid(True, alpha=0.3)

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.05, right=0.98, top=0.92, bottom=0.08, hspace=0.3, wspace=0.3)

    def update(frame):
        """Update all snapshot panels."""
        # Calculate which snapshots to show
        total_frames = len(snapshots)
        indices = [int(i * (total_frames - 1) / (n_snapshots - 1)) for i in range(n_snapshots)]

        # Update each panel
        for idx, obj, txt in zip(indices, plot_objects, time_texts):
            snap = snapshots[idx]
            t = times[idx]

            if ndim == 1:
                field = np.asarray(snap.field)
                obj.set_data(x, field)
            elif ndim == 2:
                field = np.asarray(snap.field)
                obj.set_array(field.T)
            elif ndim == 3:
                field = np.asarray(snap.field)
                obj.set_array(field[:, :, mid_z].T)

            txt.set_text(f"t = {t:.2f}")

        # Update energy plot
        if ax_energy is not None:
            current_idx = indices[-1]  # Track the last snapshot's energy
            line_E.set_data(times[: current_idx + 1], E_total[: current_idx + 1])
            energy_marker.set_data([times[current_idx]], [E_total[current_idx]])

        return plot_objects + time_texts + ([line_E, energy_marker] if ax_energy else [])

    # Create animation
    anim = FuncAnimation(
        fig, update, frames=len(snapshots), interval=interval, blit=False, repeat=True
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=1000 // interval)

    return anim

create_comparison_animation(times_list, snapshots_list, labels, potential=None, interval=50, save_path=None, figsize=(16, 5), title='Field Comparison')

Create side-by-side comparison animation for multiple simulations.

Useful for comparing different initial conditions, potentials, or parameters.

Parameters:

Name	Type	Description	Default
times_list	list[Array]	List of time arrays for each simulation	required
snapshots_list	list[list[Lattice]]	List of snapshot lists for each simulation	required
labels	list[str]	Labels for each simulation	required
potential	Callable | None	Potential function (for energy computation)	None
interval	int	Delay between frames in milliseconds	50
save_path	str | None	If provided, save animation to this path	None
figsize	tuple[int, int]	Figure size	(16, 5)
title	str	Animation title	'Field Comparison'

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Source code in jaxlatt/animation/comparison.py

def create_comparison_animation(
    times_list: list[Array],
    snapshots_list: list[list[Lattice]],
    labels: list[str],
    potential: Callable | None = None,
    interval: int = 50,
    save_path: str | None = None,
    figsize: tuple[int, int] = (16, 5),
    title: str = "Field Comparison",
) -> FuncAnimation:
    """
    Create side-by-side comparison animation for multiple simulations.

    Useful for comparing different initial conditions, potentials, or parameters.

    Args:
        times_list: List of time arrays for each simulation
        snapshots_list: List of snapshot lists for each simulation
        labels: Labels for each simulation
        potential: Potential function (for energy computation)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title

    Returns:
        FuncAnimation object
    """
    n_sims = len(snapshots_list)

    if len(times_list) != n_sims or len(labels) != n_sims:
        raise ValueError("times_list, snapshots_list, and labels must have same length")

    # Verify all are 1D
    for snapshots in snapshots_list:
        if snapshots[0].ndim != 1:
            raise ValueError("Comparison animation currently only supports 1D fields")

    # Setup figure
    fig, axes = plt.subplots(1, n_sims, figsize=figsize, sharey=True)
    if n_sims == 1:
        axes = [axes]

    # Get spatial grid (assume all same)
    lattice = snapshots_list[0][0]
    length = lattice.length[0] if isinstance(lattice.length, tuple) else lattice.length
    size = lattice.size[0] if isinstance(lattice.size, tuple) else lattice.size
    x = np.linspace(0, float(length), int(size), endpoint=False)

    # Determine y-limits
    all_fields = []
    for snapshots in snapshots_list:
        all_fields.extend([np.asarray(snap.field) for snap in snapshots])
    field_min = min(f.min() for f in all_fields)
    field_max = max(f.max() for f in all_fields)
    margin = 0.1 * (field_max - field_min)
    ylim = (field_min - margin, field_max + margin)

    # Initialize plots
    lines = []
    time_texts = []

    for i, (ax, label) in enumerate(zip(axes, labels)):
        (line,) = ax.plot([], [], "b-", linewidth=2)
        lines.append(line)

        ax.set_xlim(0, float(length))
        ax.set_ylim(ylim)
        ax.set_xlabel("x", fontsize=11)
        if i == 0:
            ax.set_ylabel("φ(x)", fontsize=11)
        ax.set_title(label, fontsize=12)
        ax.grid(True, alpha=0.3)

        time_text = ax.text(
            0.02,
            0.95,
            "",
            transform=ax.transAxes,
            fontsize=10,
            verticalalignment="top",
            bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
        )
        time_texts.append(time_text)

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.08, right=0.98, top=0.92, bottom=0.08)

    def init():
        """Initialize animation."""
        for line, time_text in zip(lines, time_texts):
            line.set_data([], [])
            time_text.set_text("")
        return lines + time_texts

    def update(frame):
        """Update animation frame."""
        for line, time_text, snapshots, times in zip(lines, time_texts, snapshots_list, times_list):
            if frame < len(snapshots):
                snap = snapshots[frame]
                t = times[frame]
                field = np.asarray(snap.field)
                line.set_data(x, field)
                time_text.set_text(f"t = {t:.3f}")

        return lines + time_texts

    # Use minimum number of frames across all simulations
    max_frames = max(len(snapshots) for snapshots in snapshots_list)

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=max_frames,
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=1000 // interval)

    return anim

field_1d(times, snapshots, potential=None, interval=50, save_path=None, figsize=(14, 5), show_energy=True, title='1D Scalar Field Evolution', ylim=None)

Create an animation of 1D scalar field evolution.

Parameters:

Name	Type	Description	Default
times	Array	Array of snapshot times	required
snapshots	list[Lattice]	List of Lattice objects at each snapshot time	required
potential	Callable | None	Potential function (needed if show_energy=True)	None
interval	int	Delay between frames in milliseconds	50
save_path	str | None	If provided, save animation to this path (e.g., 'anim.gif')	None
figsize	tuple[int, int]	Figure size (width, height)	(14, 5)
show_energy	bool	Whether to show energy evolution subplot	True
title	str	Animation title	'1D Scalar Field Evolution'
ylim	tuple[float, float] | None	Y-axis limits for field plot (auto if None)	None

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Example

from jaxlatt.potentials import quadratic_potential
anim = animate_1d_field(times, snapshots, quadratic_potential(m=1.0))
# In Jupyter: from IPython.display import HTML; HTML(anim.to_jshtml())

Source code in jaxlatt/animation/dim1.py

def field(
    times: Array,
    snapshots: list[Lattice],
    potential: Callable | None = None,
    interval: int = 50,
    save_path: str | None = None,
    figsize: tuple[int, int] = (14, 5),
    show_energy: bool = True,
    title: str = "1D Scalar Field Evolution",
    ylim: tuple[float, float] | None = None,
) -> FuncAnimation:
    """
    Create an animation of 1D scalar field evolution.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects at each snapshot time
        potential: Potential function (needed if show_energy=True)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path (e.g., 'anim.gif')
        figsize: Figure size (width, height)
        show_energy: Whether to show energy evolution subplot
        title: Animation title
        ylim: Y-axis limits for field plot (auto if None)

    Returns:
        FuncAnimation object

    Example:
        ```python
        from jaxlatt.potentials import quadratic_potential
        anim = animate_1d_field(times, snapshots, quadratic_potential(m=1.0))
        # In Jupyter: from IPython.display import HTML; HTML(anim.to_jshtml())
        ```
    """
    # Validate inputs
    if show_energy and potential is None:
        raise ValueError("Must provide potential function if show_energy=True")

    lattice = snapshots[0]
    length = lattice.length[0] if isinstance(lattice.length, tuple) else lattice.length
    size = lattice.size[0] if isinstance(lattice.size, tuple) else lattice.size
    x = np.linspace(0, float(length), int(size), endpoint=False)

    # Predefine energy arrays/handles for type checking
    E_total = E_kinetic = E_gradient = E_potential = []  # type: ignore
    line_E_total = line_E_kin = line_E_grad = line_E_pot = energy_marker = None  # type: ignore
    ax_energy = None  # type: ignore

    if show_energy:
        assert potential is not None
        energy_data = [energy_components_integrated(snap, potential) for snap in snapshots]
        E_total = [e["total"] for e in energy_data]
        E_kinetic = [e["kinetic"] for e in energy_data]
        E_gradient = [e["gradient"] for e in energy_data]
        E_potential = [e["potential"] for e in energy_data]

    # Determine y-limits for field
    if ylim is None:
        all_fields = [np.asarray(snap.field) for snap in snapshots]
        fmin = min(f.min() for f in all_fields)
        fmax = max(f.max() for f in all_fields)
        margin = 0.1 * (fmax - fmin)
        ylim = (fmin - margin, fmax + margin)

    fig = plt.figure(figsize=figsize)
    if show_energy:
        gs = GridSpec(1, 2, width_ratios=[2, 1], figure=fig)
        ax_field = fig.add_subplot(gs[0])
        ax_energy = fig.add_subplot(gs[1])  # type: ignore
    else:
        ax_field = fig.add_subplot(111)

    (line_field,) = ax_field.plot([], [], "b-", linewidth=2, label="φ(x)")
    (line_velocity,) = ax_field.plot([], [], "r--", linewidth=1.5, alpha=0.7, label="φ̇(x)")
    ax_field.set_xlim(0, float(length))
    ax_field.set_ylim(ylim)
    ax_field.set_xlabel("x", fontsize=11)
    ax_field.set_ylabel("Field value", fontsize=11)
    ax_field.grid(True, alpha=0.3)
    ax_field.legend(loc="upper right")
    time_text = ax_field.text(
        0.02,
        0.95,
        "",
        transform=ax_field.transAxes,
        fontsize=11,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    if show_energy:
        (line_E_total,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total")  # type: ignore
        (line_E_kin,) = ax_energy.plot([], [], "r-", linewidth=1.5, alpha=0.7, label="Kinetic")  # type: ignore
        (line_E_grad,) = ax_energy.plot([], [], "g-", linewidth=1.5, alpha=0.7, label="Gradient")  # type: ignore
        (line_E_pot,) = ax_energy.plot([], [], "b-", linewidth=1.5, alpha=0.7, label="Potential")  # type: ignore
        ax_energy.set_xlim(float(times[0]), float(times[-1]))  # type: ignore
        Emin = min(min(E_kinetic), min(E_gradient), min(E_potential)) * 0.9
        Emax = max(E_total) * 1.1
        ax_energy.set_ylim(Emin, Emax)  # type: ignore
        ax_energy.set_xlabel("Time", fontsize=11)  # type: ignore
        ax_energy.set_ylabel("Energy", fontsize=11)  # type: ignore
        ax_energy.legend(loc="upper right", fontsize=9)  # type: ignore
        ax_energy.grid(True, alpha=0.3)  # type: ignore
        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)  # type: ignore

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.1, right=0.95, top=0.92, bottom=0.1)

    def init():
        """Initialize animation."""
        line_field.set_data([], [])
        line_velocity.set_data([], [])
        time_text.set_text("")

        if show_energy:
            line_E_total.set_data([], [])
            line_E_kin.set_data([], [])
            line_E_grad.set_data([], [])
            line_E_pot.set_data([], [])
            energy_marker.set_data([], [])
            return (
                line_field,
                line_velocity,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return line_field, line_velocity, time_text

    def update(frame):
        """Update animation frame."""
        snap = snapshots[frame]
        t = times[frame]

        # Update field plot
        field = np.asarray(snap.field)
        field_dot = np.asarray(snap.field_dot)
        line_field.set_data(x, field)
        line_velocity.set_data(x, field_dot)
        time_text.set_text(f"t = {t:.3f}")

        # Update energy plot if needed
        if show_energy:
            line_E_total.set_data(times[: frame + 1], E_total[: frame + 1])
            line_E_kin.set_data(times[: frame + 1], E_kinetic[: frame + 1])
            line_E_grad.set_data(times[: frame + 1], E_gradient[: frame + 1])
            line_E_pot.set_data(times[: frame + 1], E_potential[: frame + 1])
            energy_marker.set_data([t], [E_total[frame]])

            return (
                line_field,
                line_velocity,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return line_field, line_velocity, time_text

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=len(snapshots),
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=1000 // interval)

    return anim

phase_space_1d(times, snapshots, x_index=None, interval=50, save_path=None, figsize=(10, 8), title='Phase Space Evolution')

Create a phase space animation showing φ vs φ̇ at a specific spatial point.

Useful for visualizing oscillatory dynamics and energy exchange.

Parameters:

Name	Type	Description	Default
times	Array	Array of snapshot times	required
snapshots	list[Lattice]	List of Lattice objects	required
x_index	int | None	Spatial index to track (center if None)	None
interval	int	Delay between frames in milliseconds	50
save_path	str | None	If provided, save animation to this path	None
figsize	tuple[int, int]	Figure size	(10, 8)
title	str	Animation title	'Phase Space Evolution'

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Source code in jaxlatt/animation/dim1.py

def phase_space(
    times: Array,
    snapshots: list[Lattice],
    x_index: int | None = None,
    interval: int = 50,
    save_path: str | None = None,
    figsize: tuple[int, int] = (10, 8),
    title: str = "Phase Space Evolution",
) -> FuncAnimation:
    """
    Create a phase space animation showing φ vs φ̇ at a specific spatial point.

    Useful for visualizing oscillatory dynamics and energy exchange.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects
        x_index: Spatial index to track (center if None)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title

    Returns:
        FuncAnimation object
    """
    lattice = snapshots[0]

    # Default to center of lattice
    if x_index is None:
        x_index = lattice.size[0] // 2

    # Extract trajectory in phase space
    field_vals = [float(np.asarray(snap.field).flatten()[x_index]) for snap in snapshots]
    field_dot_vals = [float(np.asarray(snap.field_dot).flatten()[x_index]) for snap in snapshots]

    # Setup figure
    fig, (ax_phase, ax_time) = plt.subplots(2, 1, figsize=figsize)

    # Phase space plot
    (line_phase,) = ax_phase.plot([], [], "b-", linewidth=2, alpha=0.7)
    (point_phase,) = ax_phase.plot([], [], "ro", markersize=10)

    phi_min, phi_max = min(field_vals), max(field_vals)
    phi_dot_min, phi_dot_max = min(field_dot_vals), max(field_dot_vals)
    phi_margin = 0.1 * (phi_max - phi_min)
    phi_dot_margin = 0.1 * (phi_dot_max - phi_dot_min)

    ax_phase.set_xlim(phi_min - phi_margin, phi_max + phi_margin)
    ax_phase.set_ylim(phi_dot_min - phi_dot_margin, phi_dot_max + phi_dot_margin)
    ax_phase.set_xlabel("φ", fontsize=12)
    ax_phase.set_ylabel("φ̇", fontsize=12)
    ax_phase.set_title("Phase Space Trajectory", fontsize=12)
    ax_phase.grid(True, alpha=0.3)
    ax_phase.axhline(y=0, color="k", linewidth=0.5, alpha=0.5)
    ax_phase.axvline(x=0, color="k", linewidth=0.5, alpha=0.5)

    # Time evolution plot
    (line_phi,) = ax_time.plot([], [], "b-", linewidth=2, label="φ(t)")
    (line_phi_dot,) = ax_time.plot([], [], "r-", linewidth=2, label="φ̇(t)")
    (marker_current,) = ax_time.plot([], [], "ko", markersize=8)

    ax_time.set_xlim(times[0], times[-1])
    y_min = min(phi_min, phi_dot_min) - max(phi_margin, phi_dot_margin)
    y_max = max(phi_max, phi_dot_max) + max(phi_margin, phi_dot_margin)
    ax_time.set_ylim(y_min, y_max)
    ax_time.set_xlabel("Time", fontsize=12)
    ax_time.set_ylabel("Value", fontsize=12)
    ax_time.legend(loc="upper right")
    ax_time.grid(True, alpha=0.3)
    ax_time.axhline(y=0, color="k", linewidth=0.5, alpha=0.5)

    time_text = ax_time.text(
        0.02,
        0.98,
        "",
        transform=ax_time.transAxes,
        fontsize=11,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    fig.suptitle(f"{title} (x_index = {x_index})", fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.12, right=0.95, top=0.92, bottom=0.12)

    def init():
        """Initialize animation."""
        line_phase.set_data([], [])
        point_phase.set_data([], [])
        line_phi.set_data([], [])
        line_phi_dot.set_data([], [])
        marker_current.set_data([], [])
        time_text.set_text("")
        return (
            line_phase,
            point_phase,
            line_phi,
            line_phi_dot,
            marker_current,
            time_text,
        )

    def update(frame):
        """Update animation frame."""
        t = times[frame]

        # Update phase space trajectory
        line_phase.set_data(field_vals[: frame + 1], field_dot_vals[: frame + 1])
        point_phase.set_data([field_vals[frame]], [field_dot_vals[frame]])

        # Update time series
        line_phi.set_data(times[: frame + 1], field_vals[: frame + 1])
        line_phi_dot.set_data(times[: frame + 1], field_dot_vals[: frame + 1])
        marker_current.set_data([t], [field_vals[frame]])

        time_text.set_text(f"t = {t:.3f}")

        return (
            line_phase,
            point_phase,
            line_phi,
            line_phi_dot,
            marker_current,
            time_text,
        )

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=len(snapshots),
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=1000 // interval)

    return anim

field_2d(times, snapshots, potential=None, interval=50, save_path=None, figsize=(5, 5), show_energy=False, title=None, cmap='RdBu_r', vmin=None, vmax=None, dark=False)

Create an animation of 2D scalar field evolution.

Parameters:

Name	Type	Description	Default
times	Array	Array of snapshot times	required
snapshots	list[Lattice]	List of Lattice objects at each snapshot time	required
potential	Callable | None	Potential function (needed if show_energy=True)	None
interval	int	Delay between frames in milliseconds	50
save_path	str | None	If provided, save animation to this path	None
figsize	tuple[int, int]	Figure size (width, height)	(5, 5)
show_energy	bool	Whether to show energy evolution subplot	False
title	str | None	Animation title	None
cmap	str	Colormap for field visualization	'RdBu_r'
vmin	float | None	Minimum value for colormap (auto if None)	None
vmax	float | None	Maximum value for colormap (auto if None)	None
dark	bool	Whether to use dark mode styling (white text/labels)	False

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Source code in jaxlatt/animation/dim2.py

def field(
    times: Array,
    snapshots: list[Lattice],
    potential: Callable | None = None,
    interval: int = 50,
    save_path: str | None = None,
    figsize: tuple[int, int] = (5, 5),
    show_energy: bool = False,
    title: str | None = None,
    cmap: str = "RdBu_r",
    vmin: float | None = None,
    vmax: float | None = None,
    dark: bool = False,
) -> FuncAnimation:
    """
    Create an animation of 2D scalar field evolution.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects at each snapshot time
        potential: Potential function (needed if show_energy=True)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size (width, height)
        show_energy: Whether to show energy evolution subplot
        title: Animation title
        cmap: Colormap for field visualization
        vmin: Minimum value for colormap (auto if None)
        vmax: Maximum value for colormap (auto if None)
        dark: Whether to use dark mode styling (white text/labels)

    Returns:
        FuncAnimation object
    """
    import jaxlatt.plotting as plot

    # Validate inputs
    if show_energy and potential is None:
        raise ValueError("Must provide potential function if show_energy=True")

    # Compute energy if needed
    if show_energy:
        energy_data = [energy_components_integrated(snap, potential) for snap in snapshots]
        E_total = [e["total"] for e in energy_data]
        E_kinetic = [e["kinetic"] for e in energy_data]
        E_gradient = [e["gradient"] for e in energy_data]
        E_potential = [e["potential"] for e in energy_data]

    # Determine colormap limits
    if vmin is None or vmax is None:
        all_fields = [np.asarray(snap.field) for snap in snapshots]
        field_min = min(f.min() for f in all_fields)
        field_max = max(f.max() for f in all_fields)
        if vmin is None:
            vmin = field_min
        if vmax is None:
            vmax = field_max

    # Create figure
    if show_energy:
        fig = plt.figure(figsize=figsize)
        gs = GridSpec(1, 2, width_ratios=[1, 1], figure=fig)
        ax_field = fig.add_subplot(gs[0])
        ax_energy = fig.add_subplot(gs[1])
    else:
        fig, ax_field = plt.subplots(figsize=figsize)

    # Transparent background for dark mode
    if dark:
        fig.patch.set_alpha(0.0)
        ax_field.set_facecolor("none")
        if show_energy:
            ax_energy.set_facecolor("none")

    # # Initialize field plot
    im = plot.snapshot2d(snapshots[0], ax=ax_field, vmin=vmin, vmax=vmax, cmap=cmap, dark=dark)
    time_text = ax_field.text(
        0.02,
        0.98,
        "",
        transform=ax_field.transAxes,
        fontsize=11,
        verticalalignment="top",
        color="white",
        fontweight="bold",
        bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
    )

    # Initialize energy plot if needed
    if show_energy:
        (line_E_total,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total")
        (line_E_kin,) = ax_energy.plot([], [], "r-", linewidth=1.5, alpha=0.7, label="Kinetic")
        (line_E_grad,) = ax_energy.plot([], [], "g-", linewidth=1.5, alpha=0.7, label="Gradient")
        (line_E_pot,) = ax_energy.plot([], [], "b-", linewidth=1.5, alpha=0.7, label="Potential")

        ax_energy.set_xlim(times[0], times[-1])
        E_min = min(min(E_kinetic), min(E_gradient), min(E_potential)) * 0.9
        E_max = max(E_total) * 1.1
        ax_energy.set_ylim(E_min, E_max)
        ax_energy.set_xlabel("Time", fontsize=11)
        ax_energy.set_ylabel("Energy", fontsize=11)
        ax_energy.legend(loc="upper right", fontsize=9)
        ax_energy.grid(True, alpha=0.3)

        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)

        if dark:
            plot.apply_dark_style(ax_energy)

    title_color = "white" if dark else "black"
    fig.suptitle(title, fontsize=14, fontweight="bold", color=title_color)
    # plt.subplots_adjust(left=0.1, right=0.95, top=0.92, bottom=0.1)

    def init():
        """Initialize animation."""
        im.set_array(np.asarray(snapshots[0].field).T)
        time_text.set_text("")

        if show_energy:
            line_E_total.set_data([], [])
            line_E_kin.set_data([], [])
            line_E_grad.set_data([], [])
            line_E_pot.set_data([], [])
            energy_marker.set_data([], [])
            return (
                im,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return im, time_text

    def update(frame):
        """Update animation frame."""
        snap = snapshots[frame]
        t = times[frame]

        # Update field
        field = np.asarray(snap.field)
        im.set_array(field.T)
        time_text.set_text(f"t = {t:.3f}")

        # Update energy plot if needed
        if show_energy:
            line_E_total.set_data(times[: frame + 1], E_total[: frame + 1])
            line_E_kin.set_data(times[: frame + 1], E_kinetic[: frame + 1])
            line_E_grad.set_data(times[: frame + 1], E_gradient[: frame + 1])
            line_E_pot.set_data(times[: frame + 1], E_potential[: frame + 1])
            energy_marker.set_data([t], [E_total[frame]])

            return (
                im,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return im, time_text

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=len(snapshots),
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        dpi = 150 if dark else 100
        savefig_kwargs = {"transparent": True} if dark else {}
        _save_animation(
            anim, save_path, fps=1000 // interval, dpi=dpi, savefig_kwargs=savefig_kwargs
        )

    return anim

field_3d(times, snapshots, potential=None, interval=100, save_path=None, figsize=(10, 8), title='3D Scalar Field Evolution', isosurface_levels=None, fps=20)

Create a 3D isosurface animation of scalar field evolution.

This function creates volumetric visualizations showing isosurfaces of constant field value, which is ideal for visualizing topological structures like bubbles, domain walls, and defects in 3D field configurations.

Parameters:

Name	Type	Description	Default
times	Array	Array of snapshot times	required
snapshots	list[Lattice]	List of Lattice objects (must be 3D)	required
potential	Callable | None	Potential function (optional)	None
interval	int	Delay between frames in milliseconds	100
save_path	str | None	If provided, save animation to this path	None
figsize	tuple[int, int]	Figure size (width, height)	(10, 8)
title	str	Animation title	'3D Scalar Field Evolution'
isosurface_levels	list[float] | None	List of field values to show as isosurfaces	None
fps	int	Frames per second for saved animation	20

Returns:

Type	Description
FuncAnimation	FuncAnimation object

Note

Requires matplotlib with 3D support. Isosurface rendering uses contour slices with moderate detail for performance.

Source code in jaxlatt/animation/dim3.py

def field(
    times: Array,
    snapshots: list[Lattice],
    potential: Callable | None = None,
    interval: int = 100,
    save_path: str | None = None,
    figsize: tuple[int, int] = (10, 8),
    title: str = "3D Scalar Field Evolution",
    isosurface_levels: list[float] | None = None,
    fps: int = 20,
) -> FuncAnimation:
    """
    Create a 3D isosurface animation of scalar field evolution.

    This function creates volumetric visualizations showing isosurfaces of constant
    field value, which is ideal for visualizing topological structures like bubbles,
    domain walls, and defects in 3D field configurations.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects (must be 3D)
        potential: Potential function (optional)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size (width, height)
        title: Animation title
        isosurface_levels: List of field values to show as isosurfaces
        fps: Frames per second for saved animation

    Returns:
        FuncAnimation object

    Note:
        Requires matplotlib with 3D support. Isosurface rendering uses
        contour slices with moderate detail for performance.
    """

    lattice = snapshots[0]
    if lattice.ndim != 3:
        raise ValueError("animate_3d_field requires 3D lattice")

    # Default isosurface levels
    if isosurface_levels is None:
        # Use ±1 sigma from mean as default
        all_fields = [np.asarray(snap.field) for snap in snapshots]
        field_mean = np.mean([f.mean() for f in all_fields])
        field_std = np.mean([f.std() for f in all_fields])
        isosurface_levels = [field_mean + field_std, field_mean - field_std]

    nx, ny, nz = lattice.size
    lx, ly, lz = lattice.length

    # Create coordinate grids
    x = np.linspace(0, lx, nx)
    y = np.linspace(0, ly, ny)
    z = np.linspace(0, lz, nz)
    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")

    # Setup figure
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, projection="3d")

    # Set viewing angle
    ax.view_init(elev=20, azim=45)

    # Set axis limits
    ax.set_xlim(0, lx)
    ax.set_ylim(0, ly)
    ax.set_zlim(0, lz)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")

    time_text = fig.text(
        0.02,
        0.95,
        "",
        fontsize=12,
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.7),
    )

    fig.suptitle(title, fontsize=14, fontweight="bold")

    def update(frame):
        """Update animation frame."""
        ax.clear()

        snap = snapshots[frame]
        t = times[frame]
        field = np.asarray(snap.field)

        # Set axis properties again after clear
        ax.set_xlim(0, lx)
        ax.set_ylim(0, ly)
        ax.set_zlim(0, lz)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("z")
        ax.view_init(elev=20, azim=45 + frame * 0.5)  # Slow rotation

        # Plot isosurfaces using contour slices
        # We'll show XY, XZ, and YZ slices with contours
        colors = ["blue", "red", "green", "orange"]

        for idx, level in enumerate(isosurface_levels[:4]):  # Max 4 levels
            color = colors[idx % len(colors)]

            # XY slices (at various Z positions)
            for zi in range(0, nz, max(1, nz // 5)):
                field_slice = field[:, :, zi]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        X[:, :, zi],
                        Y[:, :, zi],
                        field_slice,
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=z[zi],
                        zdir="z",
                    )

            # XZ slices (at various Y positions)
            for yi in range(0, ny, max(1, ny // 5)):
                field_slice = field[:, yi, :]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        X[:, yi, :],
                        field_slice,
                        Z[:, yi, :],
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=y[yi],
                        zdir="y",
                    )

            # YZ slices (at various X positions)
            for xi in range(0, nx, max(1, nx // 5)):
                field_slice = field[xi, :, :]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        Y[xi, :, :],
                        field_slice,
                        Z[xi, :, :],
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=x[xi],
                        zdir="x",
                    )

        time_text.set_text(f"t = {t:.3f}")

        return ax, time_text

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        frames=len(snapshots),
        interval=interval,
        blit=False,  # 3D animations can't use blitting
        repeat=True,
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=fps)
        console.print("[green]✓ Animation saved![/green]")

    return anim

slices_3d(times, snapshots, potential=None, interval=100, save_path=None, figsize=(14, 6), title='3D Scalar Field Central Slices', cmap='RdBu_r', vmin=None, vmax=None)

Animate central XY, XZ, YZ slices of a 3D scalar field.

Parameters:

Name	Type	Description	Default
times	Array	Snapshot times	required
snapshots	list[Lattice]	Lattice states (ndim must be 3)	required
potential	Callable | None	Potential function (optional; if provided energy displayed)	None
interval	int	Frame delay ms	100
save_path	str | None	Optional path to save animation (GIF or MP4)	None
figsize	tuple[int, int]	Figure size	(14, 6)
title	str	Overall title	'3D Scalar Field Central Slices'
cmap	str	Colormap	'RdBu_r'
vmin/vmax		Explicit color scale limits	required

Returns:

Type	Description
FuncAnimation	FuncAnimation

Source code in jaxlatt/animation/dim3.py

def slices(
    times: Array,
    snapshots: list[Lattice],
    potential: Callable | None = None,
    interval: int = 100,
    save_path: str | None = None,
    figsize: tuple[int, int] = (14, 6),
    title: str = "3D Scalar Field Central Slices",
    cmap: str = "RdBu_r",
    vmin: float | None = None,
    vmax: float | None = None,
) -> FuncAnimation:
    """
    Animate central XY, XZ, YZ slices of a 3D scalar field.

    Args:
        times: Snapshot times
        snapshots: Lattice states (ndim must be 3)
        potential: Potential function (optional; if provided energy displayed)
        interval: Frame delay ms
        save_path: Optional path to save animation (GIF or MP4)
        figsize: Figure size
        title: Overall title
        cmap: Colormap
        vmin/vmax: Explicit color scale limits

    Returns:
        FuncAnimation
    """
    lattice0 = snapshots[0]
    if lattice0.ndim != 3:
        raise ValueError("animate_3d_slices requires 3D lattice snapshots")

    # Determine color limits
    if vmin is None or vmax is None:
        all_fields = [np.asarray(s.field) for s in snapshots]
        fmin = min(f.min() for f in all_fields)
        fmax = max(f.max() for f in all_fields)
        if vmin is None:
            vmin = fmin
        if vmax is None:
            vmax = fmax

    nx, ny, nz = lattice0.size
    cx, cy, cz = nx // 2, ny // 2, nz // 2

    # Prepare figure
    fig = plt.figure(figsize=figsize)
    gs = GridSpec(2, 3, figure=fig, height_ratios=[1, 0.05])
    ax_xy = fig.add_subplot(gs[0, 0])
    ax_xz = fig.add_subplot(gs[0, 1])
    ax_yz = fig.add_subplot(gs[0, 2])
    ax_cbar = fig.add_subplot(gs[1, :])

    extent_xy = (0.0, lattice0.length[0], 0.0, lattice0.length[1])
    extent_xz = (0.0, lattice0.length[0], 0.0, lattice0.length[2])
    extent_yz = (0.0, lattice0.length[1], 0.0, lattice0.length[2])

    field0 = np.asarray(lattice0.field)
    im_xy = ax_xy.imshow(
        field0[:, :, cz].T,
        origin="lower",
        cmap=cmap,
        extent=extent_xy,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )
    im_xz = ax_xz.imshow(
        field0[:, cy, :].T,
        origin="lower",
        cmap=cmap,
        extent=extent_xz,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )
    im_yz = ax_yz.imshow(
        field0[cx, :, :].T,
        origin="lower",
        cmap=cmap,
        extent=extent_yz,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )

    # Set titles & axis labels with correct coordinates per slice
    ax_xy.set_title("XY (z=mid)")
    ax_xy.set_xlabel("x")
    ax_xy.set_ylabel("y")

    ax_xz.set_title("XZ (y=mid)")
    ax_xz.set_xlabel("x")
    ax_xz.set_ylabel("z")

    ax_yz.set_title("YZ (x=mid)")
    ax_yz.set_xlabel("y")
    ax_yz.set_ylabel("z")

    cbar = plt.colorbar(im_xy, cax=ax_cbar, orientation="horizontal")
    cbar.set_label("φ")

    time_text = fig.text(
        0.02,
        0.95,
        "",
        fontsize=12,
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.6),
    )

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.08, right=0.95, top=0.92, bottom=0.08)

    def update(frame):
        snap = snapshots[frame]
        f = np.asarray(snap.field)
        im_xy.set_array(f[:, :, cz].T)
        im_xz.set_array(f[:, cy, :].T)
        im_yz.set_array(f[cx, :, :].T)
        time_text.set_text(f"t = {times[frame]:.3f}")
        return im_xy, im_xz, im_yz, time_text

    anim = FuncAnimation(
        fig,
        update,
        frames=len(snapshots),
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        _save_animation(anim, save_path, fps=1000 // interval)

    return anim