Visualization

Install the optional matplotlib backend before calling these functions:

pip install 'jax-morph[viz]'

With uv, use uv add 'jax-morph[viz]'. The base package keeps matplotlib optional, while import jax_morph.viz remains safe without the extra.

All functions return native matplotlib objects, so normal axes customization, notebook display, and animation export remain available to the caller. animate(..., ax=ax) can target an existing compatible axes; frame updates replace only the animation's own cell artist and categorical legend, leaving caller-added artists, labels, and annotations intact.

Static rendering#

jax_morph.viz.draw #

draw(
    state,
    *,
    color_by=None,
    cmap='viridis',
    vlim=None,
    categorical=None,
    color='tab:blue',
    colorbar=True,
    alpha=0.85,
    render_3d='spheres',
    sphere_resolution=None,
    marker_scale=None,
    ax=None,
    lims=None,
    grid=False,
    colorbar_label=None,
    title='',
)

Draw the live cells of one unstacked state on a native matplotlib axes.

One-dimensional states are drawn as a circle strip, while two-dimensional circles and three- dimensional sphere meshes use physical data-space radii. In 3-D, render_3d='markers' is a faster screen-space approximation: marker area scales with radius**2 but absolute size, overlap, and perspective are not physical. marker_scale tunes that approximation.

Parameters:

state –

Unstacked simulation state to render.
color_by –

Optional scalar-per-cell field name, callable, or raw array.
cmap –

Matplotlib colormap for selected values. Defaults to 'viridis'.
vlim –

Optional continuous (low, high) color range. Defaults to None.
categorical –

Whether to render selected values as discrete categories. Defaults to None.
color –

Solid cell color when color_by is None. Defaults to 'tab:blue'.
colorbar –

Whether to add a continuous colorbar. Defaults to True.
alpha –

Cell opacity. Defaults to 0.85.
render_3d –

'spheres' for physical meshes or 'markers' for fast markers.
sphere_resolution –

Icosphere subdivision level from 0 through 3. Defaults to 3 for spheres.
marker_scale –

Points-squared multiplier for 3-D markers. Defaults to 1800.0 for markers.
ax –

Optional compatible matplotlib axes. Defaults to None.
lims –

Optional physical data limits, one pair per spatial dimension. Defaults to None.
grid –

Whether to retain contextual ticks, panes, and a light grid. Defaults to False.
colorbar_label –

Optional continuous colorbar label. Named fields label themselves by default.
title –

Axes title. Defaults to an empty string.

Returns:

–

The matplotlib axes containing the cell artists.

3-D sphere cost.

The sphere default sphere_resolution=3 favors fidelity and emits 1,280 depth-sorted faces per cell, so a large cluster can be slow to draw. Lower sphere_resolution or use render_3d='markers' for quick iteration on many cells.

Animation#

jax_morph.viz.animate #

animate(
    history,
    *,
    color_by=None,
    cmap='viridis',
    vlim=None,
    categorical=None,
    color='tab:blue',
    colorbar=True,
    alpha=0.85,
    render_3d='markers',
    sphere_resolution=None,
    marker_scale=None,
    interval=300,
    fixed_lims=True,
    lims=None,
    fixed_clim=True,
    fig=None,
    ax=None,
    grid=False,
    colorbar_label=None,
    title=None,
)

Animate a complete stacked history or TrajectoryRecord with matplotlib.

The function replaces only its own cell artist and categorical legend as cells appear, disappear, move, and change radius; other artists and axes labels remain intact. Three- dimensional markers are the default for predictable iteration speed; opt into render_3d='spheres' for physical data-space radii. The returned animation controls saving and notebook embedding.

Parameters:

history –

Stacked complete state history or trajectory record.
color_by –

Optional scalar-per-cell field name, callable, or raw array.
cmap –

Matplotlib colormap for selected values. Defaults to 'viridis'.
vlim –

Optional continuous (low, high) color range. Defaults to None.
categorical –

Whether to render selected values as discrete categories. Defaults to None.
color –

Solid cell color when color_by is None. Defaults to 'tab:blue'.
colorbar –

Whether to add one stable continuous colorbar. Defaults to True.
alpha –

Cell opacity. Defaults to 0.85.
render_3d –

'spheres' or 'markers'. Defaults to 'markers'.
sphere_resolution –

Icosphere level from 0 through 3. Defaults to 0 for spheres.
marker_scale –

Points-squared multiplier for markers. Defaults to 1800.0.
interval –

Delay between frames in milliseconds. Defaults to 300.
fixed_lims –

Whether to keep one view across all frames. Defaults to True.
lims –

Optional persistent physical data limits, one pair per spatial dimension. Requires fixed_lims=True. Defaults to None.
fixed_clim –

Whether to keep one continuous range across all frames. Defaults to True.
fig –

Optional empty matplotlib figure to populate. Mutually exclusive with ax. Defaults to None.
ax –

Optional compatible matplotlib axes to populate. Mutually exclusive with fig. Defaults to None.
grid –

Whether to retain contextual ticks, panes, and a light grid. Defaults to False.
colorbar_label –

Optional continuous colorbar label. Named fields label themselves by default.
title –

Fixed title, or None to display each frame time. Defaults to None.

Returns:

–

A matplotlib FuncAnimation.

Per-cell time series#

jax_morph.viz.plot_timeseries #

plot_timeseries(
    history,
    field,
    *,
    index=None,
    cell_ids=None,
    ax=None,
    alpha=0.85,
    linewidth=1.5,
    title='',
    ylabel=None,
)

Plot one selected cell-scope value against recorded simulation time.

Cell traces use trajectory-local reconstructed ids rather than storage slots, so a slot reused after death creates a new line. Named fields and callables may have trailing components selected by index; raw arrays are already scalar per cell and may be constant or frame-aligned.

Parameters:

history –

Stacked complete state history or trajectory record.
field –

Cell field name, callable selector, or scalar raw array.
index –

Optional complete trailing component index for a named or callable field.
cell_ids –

Optional cell id or ordered iterable of ids to draw.
ax –

Optional ordinary matplotlib axes. Defaults to None.
alpha –

Line opacity. Defaults to 0.85.
linewidth –

Non-negative matplotlib line width. Defaults to 1.5.
title –

Axes title. Defaults to an empty string.
ylabel –

Optional y-axis label overriding the selected-field default.

Returns:

–

The matplotlib axes containing one line per selected cell identity.