Euler angle playground widget for napari

euler_playground.py

#!/usr/bin/env python3

import requests
from io import StringIO
from functools import lru_cache

import numpy as np
from scipy.spatial.transform import Rotation, Slerp
import napari
from napari.settings import get_settings
from napari.utils.notifications import show_info
from magicgui import magicgui


@lru_cache
def teapot():
    # download the utah teapot model
    url = 'https://users.cs.utah.edu/~dejohnso/models/teapot_bezier0.tris'
    response = requests.get(url)
    if response.status_code != 200:
        return None
    # parse it as numpy array
    vert = np.loadtxt(StringIO(response.text.partition('\n')[2]))
    # rotate and rescale so it's along z in napari
    vert = Rotation.from_euler('z', -90, degrees=True).apply(vert) * 2
    faces = np.arange(len(vert)).reshape(-1, 3)
    return vert, faces


@magicgui(
    auto_call=True,
    rot=dict(widget_type="Slider", min=0, max=360),
    tilt=dict(widget_type="Slider", min=0, max=360),
    psi=dict(widget_type="Slider", min=0, max=360),
    mode=dict(choices=["extrinsic", "intrinsic"]),
    handedness=dict(choices=["right", "left"]),
)
def euler(
    viewer: napari.Viewer,
    rot,
    tilt,
    psi,
    axes="zyz",
    mode="extrinsic",
    handedness="right",
) -> napari.types.LayerDataTuple:
    """Euler angle playground.

    Simple widget to generate rotations from euler angles with different
    combinations of angles and conventions.

    Also generates a small time-series showcasing the rotation from the original
    frame of reference to the new one.
    """

    if not len(axes) == 3 or any(x not in "xyzXYZ" for x in axes):
        show_info(
            f"{axes} is not a valid combination of axes (3 axes among x, y and z)"
        )
        return []

    # Rotation object uses case to distinguish convention
    axes = axes.upper() if mode == "intrinsic" else axes.lower()

    if handedness == "left":
        rot, tilt, psi = -rot, -tilt, -psi

    angles = [rot, tilt, psi]

    rotation = Rotation.from_euler(axes, angles, degrees=True)

    timeseries_length = 60

    # prepare empty 4D array formatted for napari vectors
    # N * 3 axes, origin+direction, xyz
    vectors_napari = np.zeros((timeseries_length * 3, 2, 4))
    # set increasing origins along time dimension
    vectors_napari[:, 0, 0] = np.arange(timeseries_length).repeat(3)

    # interpolate between the original reference frame and the final rotation
    sl = Slerp([0, 1], Rotation.concatenate([Rotation.identity(), rotation]))
    rot_interp = sl(np.linspace(0, 1, timeseries_length))

    # loop over x y and z and apply the rotations to each
    for i, unit_vec in enumerate(np.eye(3)):
        rotated = rot_interp.apply(unit_vec)
        # set the value of the new basis vector (every 3 elements so we get xyzxyzxyz and so on)
        # note: we're inverting xyz to zyx cause napari is in zyx coordinate system
        vectors_napari[i::3, 1, 1:] = rotated[..., ::-1]

    # jump to last time point so we see the final rotation
    viewer.dims.set_current_step(0, viewer.dims.range[0][1])

    layers = [
        (
            vectors_napari,
            dict(
                name="rotated axes",
                edge_color=np.tile(
                    np.array(["cyan", "magenta", "yellow"]), timeseries_length
                ),
                length=10,
                vector_style="arrow",
            ),
            "vectors",
        ),
    ]

    if teapot() is not None:
        rotmat = rotation.as_matrix()[::-1, ::-1]
        layers.append(
            (
                teapot(),
                dict(
                    name='rotated object',
                    rotate=rotmat,
                ),
                "surface"
            ),
        )

    return layers


if __name__ == "__main__":
    viewer = napari.Viewer()
    viewer.window.add_dock_widget(euler, name="Euler angles playground")

    # run the widget once to generate the layer
    euler()

    # set up napari viewer
    viewer.dims.ndisplay = 3
    viewer.axes.visible = True
    viewer.dims.axis_labels = list("zyx")
    get_settings().application.playback_fps = 60
    viewer.dims.set_current_step(0, viewer.dims.range[0][1])

    # center and rotate camera so it's all nice and visible
    viewer.camera.center = (0, 0, 0)
    viewer.camera.zoom = 15
    viewer.camera.angles = (-158.037121343505, 39.971870082826875, -148.78974354446032)

    # start napari event loop
    napari.run()

To install dependency and run:

pip install scipy "napari[all]"
python euler_playground.py

euler_angles.mp4