Note

This example is available as a jupyter notebook here.

import ring
import mediapy as media
import jax
from diodem import load_data

sys_str = """
<x_xy>
<worldbody>
<geom dim="0.1" type="xyz"></geom>
<body joint="free" name="seg" pos="0 0 .5">
<geom color="dustin_exp_blue" dim="0.15 0.075 0.05" mass="0.2" pos="0.03 0 0" type="box"></geom>
<body joint="frozen" name="imu" pos="0.0 0.0 0.03">
<geom color="dustin_exp_orange" dim="0.05 0.03 0.02" mass="0.1" type="box"></geom>
</body>
</body>
</worldbody>
</x_xy>
"""

sys = ring.System.create(sys_str)

(X, y), (key, q, x, _) = ring.RCMG(sys, keep_output_extras=True).to_list()[0]

executing generators: 100%|██████████| 1/1 [00:07<00:00,  7.04s/it]


media.show_video(sys.render(xs=x, width=640, height=480, camera="target", render_every_nth=4), fps=25)

Rendering frames..: 100%|██████████| 1500/1500 [00:09<00:00, 154.24it/s]


exp_data = load_data(exp_id=1, motion_start="fast")

exp_data.keys()

dict_keys(['seg1', 'seg2', 'seg3', 'seg4', 'seg5'])

exp_data["seg1"].keys()

dict_keys(['imu_nonrigid', 'imu_rigid', 'marker1', 'marker2', 'marker3', 'marker4', 'quat'])

segment = "seg2"
omc_data_sys = {
    "seg": {
        "pos": exp_data[segment]["marker1"],
        "quat": exp_data[segment]["quat"],
    },
    "imu": {
        "quat": exp_data[segment]["quat"],
    }
}

omc_data_sys

{'seg': {'pos': array([[-0.00217298,  1.31229643, -0.06881995],
         [-0.00217508,  1.31225208, -0.06880918],
         [-0.00217637,  1.31223445, -0.06880335],
         ...,
         [-0.07191098,  1.33661952, -0.07041251],
         [-0.07191038,  1.33658217, -0.07042007],
         [-0.07190987,  1.3365283 , -0.07043372]]),
  'quat': array([[ 0.95854069,  0.11865217, -0.03095211, -0.25722244],
         [ 0.95854969,  0.11859988, -0.03091764, -0.25721714],
         [ 0.95855869,  0.11854759, -0.03088317, -0.25721185],
         ...,
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111],
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111],
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111]])},
 'imu': {'quat': array([[ 0.95854069,  0.11865217, -0.03095211, -0.25722244],
         [ 0.95854969,  0.11859988, -0.03091764, -0.25721714],
         [ 0.95855869,  0.11854759, -0.03088317, -0.25721185],
         ...,
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111],
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111],
         [ 0.95514079,  0.11290525, -0.03582784, -0.27143111]])}}

x = ring.sim2real.xs_from_raw(sys, omc_data_sys)

# vectorize this function over time
@jax.vmap
def update_position_vector_of_imu(x):
    state = ring.State.create(sys, x=x)
    # populate minimal coordinates `state.q` from maximal coordinates `state.x`
    state = ring.algorithms.inverse_kinematics(sys, state)
    # re-calculate maximal coordiantes `state.x` from minimal coordinates `state.q`
    # this uses the position vector specified in the system (and so the xml file)
    # to produce an offset between IMu and segment geom box
    _, state = ring.algorithms.forward_kinematics(sys, state)
    return state.x

x = update_position_vector_of_imu(x)

media.show_video(sys.render(xs=x, width=640, height=480, camera="target", render_every_nth=4), fps=25)

Rendering frames..: 100%|██████████| 1075/1075 [00:06<00:00, 156.26it/s]