Visualize Cylinder with Wrench

We apply a constant body-fixed wrench to a cylinder and integrate acceleration to twist and exponential coordinates of transformation to finally compute the new pose of the cylinder.

import numpy as np

import pytransform3d.visualizer as pv
from pytransform3d.transformations import transform_from_exponential_coordinates


def spatial_inertia_of_cylinder(mass, length, radius):
    I_xx = I_yy = 0.25 * mass * radius**2 + 1.0 / 12.0 * mass * length**2
    I_zz = 0.5 * mass * radius**2
    inertia = np.eye(6)
    inertia[:3, :3] *= np.array([I_xx, I_yy, I_zz])
    inertia[3:, 3:] *= mass
    return inertia


def animation_callback(
    step, cylinder, cylinder_frame, prev_cylinder2world, Stheta_dot, inertia_inv
):
    if step == 0:  # Reset cylinder state
        prev_cylinder2world[:, :] = np.eye(4)
        Stheta_dot[:] = 0.0

    # Apply constant wrench
    wrench_in_cylinder = np.array([0.1, 0.001, 0.001, 0.01, 1.0, 1.0])
    dt = 0.0005

    Stheta_ddot = np.dot(inertia_inv, wrench_in_cylinder)
    Stheta_dot += dt * Stheta_ddot
    cylinder2world = transform_from_exponential_coordinates(
        dt * Stheta_dot
    ).dot(prev_cylinder2world)

    # Update visualization
    cylinder_frame.set_data(cylinder2world)
    cylinder.set_data(cylinder2world)

    prev_cylinder2world[:, :] = cylinder2world

    return cylinder_frame, cylinder


fig = pv.figure()

# Definition of cylinder
mass = 1.0
length = 0.5
radius = 0.1
inertia_inv = np.linalg.inv(
    spatial_inertia_of_cylinder(mass=mass, length=length, radius=radius)
)

# State of cylinder
cylinder2world = np.eye(4)
twist = np.zeros(6)

cylinder = fig.plot_cylinder(length=length, radius=radius, c=[1, 0.5, 0])
cylinder_frame = fig.plot_transform(A2B=cylinder2world, s=0.5)

fig.plot_transform(A2B=np.eye(4), s=0.5)

fig.view_init()

if "__file__" in globals():
    fig.animate(
        animation_callback,
        n_frames=10000,
        fargs=(cylinder, cylinder_frame, cylinder2world, twist, inertia_inv),
        loop=True,
    )
    fig.show()
else:
    fig.save_image("__open3d_rendered_image.jpg")