Cube - Rerun — ContextQMD

Rerun imports and initialization

from __future__ import annotations

import math
import uuid
from collections import namedtuple
from math import cos, sin

import numpy as np

import rerun as rr  # pip install rerun-sdk
import rerun.blueprint as rrb
from rerun.notebook import Viewer  # pip install rerun-notebook

Helper to create the colored cube

This code exists in the rerun.utilities package, but is included here for context.

ColorGrid = namedtuple("ColorGrid", ["positions", "colors"])


def build_color_grid(x_count: int = 10, y_count: int = 10, z_count: int = 10, twist: float = 0) -> ColorGrid:
    """
    Create a cube of points with colors.

    The total point cloud will have x_count * y_count * z_count points.

    Parameters
    ----------
    x_count, y_count, z_count:
        Number of points in each dimension.
    twist:
        Angle to twist from bottom to top of the cube

    """

    grid = np.mgrid[
        slice(-x_count, x_count, x_count * 1j),
        slice(-y_count, y_count, y_count * 1j),
        slice(-z_count, z_count, z_count * 1j),
    ]

    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)
    for z in range(z_count):
        xv, yv, zv = grid[:, :, :, z]
        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])
        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])
        grid[:, :, :, z] = [rot_xv, rot_yv, zv]

    positions = np.vstack([xyz.ravel() for xyz in grid])

    colors = np.vstack([
        xyz.ravel()
        for xyz in np.mgrid[
            slice(0, 255, x_count * 1j),
            slice(0, 255, y_count * 1j),
            slice(0, 255, z_count * 1j),
        ]
    ])

    return ColorGrid(positions.T, colors.T.astype(np.uint8))

Logging some data

Now we can log some data and add it to the recording, and display it using notebook_show.

Note that displaying a recording will consume the data, so it will not be available for use in a subsequent cell.

rr.init("rerun_example_cube")

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rr.set_time("step", sequence=t)
    cube = build_color_grid(10, 10, 10, twist=twists[t])
    rr.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))

rr.notebook_show()

Logging live data

Using rr.notebook_show like above buffers the data in the recording stream, but doesn't process it until the call to rr.notebook_show.

However, rr.notebook_show can be called at any time during your cell's execution to immediately display the Rerun viewer. You can then incrementally stream to it. Here we add a sleep to simulate a cell that does a lot more processing. By calling notebook_show first, we can see the output of our code live while it's still running.

from time import sleep

rr.init("rerun_example_cube")

rr.notebook_show()

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    sleep(0.05)
    rr.set_time("step", sequence=t)
    cube = build_color_grid(10, 10, 10, twist=twists[t])
    rr.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))

Incremental logging

Note that until we either reset the recording stream (by calling rr.init()), or create a new output widget (by calling rr.notebook_show() The currently active stream in the kernel will continue to send events to the existing widget.

The following will add a rotation to the above recording.

for t in range(STEPS):
    sleep(0.01)
    rr.set_time("step", sequence=t)
    rr.log("cube", rr.Transform3D(rotation=rr.RotationAxisAngle(axis=[1, 0, 0], degrees=t)))

Starting a new recording

You can always start another recording by calling rr.init(...) again to reset the global stream, or alternatively creating a separate recording stream using the rr.RecordingStream constructor (discussed more below)

rr.init("rerun_example_cube")

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rr.set_time("step", sequence=t)
    h_grid = build_color_grid(10, 3, 3, twist=twists[t])
    rr.log("h_grid", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))
    v_grid = build_color_grid(3, 3, 10, twist=twists[t])
    rr.log("v_grid", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))

rr.notebook_show()

Adjusting the view

The notebook_show method also lets you adjust properties such as width and height.

rr.init("rerun_example_cube")

small_cube = build_color_grid(3, 3, 3, twist=0)
rr.log("small_cube", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))

rr.notebook_show(height=400)

To update the default width and height, you can use the rerun.notebook.set_default_size function.

Note that you do not need to initialize a recording to use this function.

from rerun.notebook import set_default_size

set_default_size(width=400, height=400)

rr.init("rerun_example_cube")

small_cube = build_color_grid(3, 3, 3, twist=0)
rr.log("small_cube", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))

rr.notebook_show()

set_default_size(width=640, height=480)

Using blueprints

Rerun blueprints can be used with rr.notebook_show()

For example, we can split the two grids into their own respective views.

rr.init("rerun_example_cube")

blueprint = rrb.Blueprint(
    rrb.Horizontal(
        rrb.Spatial3DView(name="Horizontal grid", origin="h_grid"),
        rrb.Spatial3DView(name="Vertical grid", origin="v_grid"),
        column_shares=[2, 1],
    ),
    collapse_panels=True,
)

rr.notebook_show(blueprint=blueprint)

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rr.set_time("step", sequence=t)
    h_grid = build_color_grid(10, 3, 3, twist=twists[t])
    rr.log("h_grid", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))
    v_grid = build_color_grid(3, 3, 10, twist=twists[t])
    rr.log("v_grid", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))

Extra convenience

Rerun blueprints types also implement _ipython_display_() directly, so if a blueprint is the last element in your cell the right thing will happen.

Note that this mechanism only works when you are using the global recording stream.

rr.init("rerun_example_cube")

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rr.set_time("step", sequence=t)
    h_grid = build_color_grid(10, 3, 3, twist=twists[t])
    rr.log("h_grid", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))
    v_grid = build_color_grid(3, 3, 10, twist=twists[t])
    rr.log("v_grid", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))

rrb.Spatial3DView(name="Horizontal grid", origin="h_grid")

Working with non-global streams

Sometimes it can be more explicit to work with specific (non-global recording) streams via rr.RecordingStream constructor.

In this case, remember to call notebook_show directly on the recording stream. As noted above, there is no way to use a bare Blueprint object in conjunction with a non-global recording.

rec = rr.RecordingStream("rerun_example_cube_flat")

bp = rrb.Blueprint(collapse_panels=True)

rec.notebook_show(blueprint=bp)

flat_grid = build_color_grid(20, 20, 1, twist=0)
rec.log("flat_grid", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))

Using the Viewer object directly

Instead of calling notebook_show you can alternatively hold onto the viewer object.

This lets you add additional recordings to a view.

rec = rr.RecordingStream("rerun_example_multi_recording", recording_id=uuid.uuid4())

flat_grid = build_color_grid(20, 20, 1, twist=0)
rec.log("flat_grid", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))

viewer = Viewer(recording=rec, blueprint=rrb.Blueprint(rrb.BlueprintPanel(state="expanded")))
viewer.display()

rec = rr.RecordingStream("rerun_example_multi_recording", recording_id=uuid.uuid4())

viewer.add_recording(rec)

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rec.set_time("step", sequence=t)
    cube = build_color_grid(10, 10, 10, twist=twists[t])
    rec.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))

Controlling the Viewer

Other than sending a blueprint to the Viewer, some parts of it can also be controlled directly through Python.

viewer = Viewer()
viewer.display()

recordings = [
    rr.RecordingStream("rerun_example_time_ctrl", recording_id="example_a"),
    rr.RecordingStream("rerun_example_time_ctrl", recording_id="example_b"),
]

rec_colors = {"example_a": [0, 255, 0], "example_b": [255, 0, 0]}

for rec in recordings:
    viewer.add_recording(rec)

STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for rec in recordings:
    for t in range(STEPS):
        cube = build_color_grid(10, 10, 10, twist=twists[t])
        rec.set_time("step", sequence=t)
        rec.log("cube", rr.Points3D(cube.positions, colors=rec_colors[rec.get_recording_id()], radii=0.5))

The state of each panel in the viewer can be overridden, locking it in the specified state.

viewer.update_panels(blueprint="expanded")

In multi-recording scenarios, the active recording can be set using set_active_recording. The timeline panel's state for the currently active recording can be controlled using set_time_ctrl.

viewer.set_active_recording(recording_id="example_a")
viewer.set_time_ctrl(timeline="step", sequence=25)
viewer.set_active_recording(recording_id="example_b")
viewer.set_time_ctrl(timeline="step", sequence=75)

Switch between the two recordings in the blueprint panel to see the updated timelines.