imt-ring 1.2.1__py3-none-any.whl

ring/algorithms/sensors.py

@@ -0,0 +1,582 @@
+from collections import defaultdict
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+from ring import algebra
+from ring import algorithms
+from ring import base
+from ring import io
+from ring import maths
+from ring import sim2real
+
+
+def accelerometer(
+    xs: base.Transform, gravity: jax.Array, dt: float, n: int
+) -> jax.Array:
+    """Compute measurements of an accelerometer that follows a frame which moves along
+    a trajectory of Transforms. Let `xs` be the transforms from base to link.
+    """
+
+    acc = (xs.pos[: -2 * n] + xs.pos[2 * n :] - 2 * xs.pos[n:-n]) / (n * dt) ** 2
+    acc = acc + gravity
+
+    # 2nd order derivative, (N,) -> (N-2n,)
+    # prepend and append n elements to keep shape size
+    acc = jnp.vstack((jnp.atleast_2d(acc[:n]), acc, jnp.atleast_2d(acc[-n:])))
+
+    return maths.rotate(acc, xs.rot)
+
+
+def gyroscope(rot: jax.Array, dt: float, second_order: bool) -> jax.Array:
+    """Compute measurements of a gyroscope that follows a frame with an orientation
+    given by trajectory of quaternions `rot`."""
+    # this was not needed before
+    # the reason is that before q represented
+    # FROM LOCAL TO EPS
+    # but now we q represents
+    # FROM EPS TO LOCAL
+    # q = maths.quat_inv(rot)
+
+    q = rot
+    if second_order:
+        dq = maths.quat_mul(q[2:], maths.quat_inv(q[:-2]))
+        dq = jnp.vstack((dq[0][None], dq, dq[-1][None]))
+
+        dt = 2 * dt
+    else:
+        # 1st-order approx to derivative
+        dq = maths.quat_mul(q[1:], maths.quat_inv(q[:-1]))
+
+        # due to 1st order derivative, shape (N,) -> (N-1,)
+        # append one element at the end to keep shape size
+        dq = jnp.vstack((dq, dq[-1][None]))
+
+    axis, angle = maths.quat_to_rot_axis(dq)
+    angle = angle[:, None]
+
+    gyr = axis * angle / dt
+    return jnp.where(jnp.abs(angle) > 1e-10, gyr, jnp.zeros(3))
+
+
+def _draw_random_magvec(key):
+    "Unit is in a.u. (40 microTesla)"
+    c1, c2 = jax.random.split(key)
+
+    dip_angle_min, dip_angle_max = -85.0, -50.0  # degrees
+    dip_angle = jnp.deg2rad(
+        jax.random.uniform(c1, minval=dip_angle_min, maxval=dip_angle_max)
+    )
+
+    norm_minval, norm_maxval = 15e-6, 65e-6  # Tesla; from lecture script page 10
+    # convert Tesla -> a.u. where (1 a.u. ~ 40 microTesla)
+    au = 40e-6
+    norm_minval, norm_maxval = norm_minval / au, norm_maxval / au
+    norm = jax.random.uniform(c2, minval=norm_minval, maxval=norm_maxval)
+
+    return jnp.array([0.0, jnp.cos(dip_angle), jnp.sin(dip_angle)]) * norm
+
+
+def magnetometer(rot: jax.Array, magvec: jax.Array) -> jax.Array:
+    return maths.rotate(magvec, rot)
+
+
+# Xsens MTI 10
+# gyr:
+# - bias error: 0.2 deg/s
+# - bias stability: 18 deg/h
+# - noise density: 0.03 deg/s/sqrt(hz)
+# -> 40 hz: 0.2 deg/s
+# -> 100 hz: 0.3 deg/s
+# acc:
+# - bias error: 0.05 m/s/s
+# - bias stability: 15 micro g (<- gravity)
+# - noise density: 60 micro g/sqrt(hz)
+# -> 40 hz: 0.0036 m/s/s
+# -> 100 hz: 0.006 m/s/s
+# mag:
+# - Total RMS noise: 0.5 milliGauss (1 Gauss = 1e-4 Tesla)
+# -------------
+# Xsens MTI 100
+# gyr:
+# - bias error: 0.2 deg/s
+# - bias stability: 10 deg/h
+# - noise density: 0.01 deg/s/sqrt(hz)
+# -> 40 hz: 0.067 deg/s
+# -> 100 hz: 0.1 deg/s
+# acc:
+# - bias error: 0.05 m/s/s
+# - bias stability: 15 micro g (<- gravity)
+# - noise density: 60 micro g/sqrt(hz)
+# -> 40 hz: 0.0036 m/s/s
+# -> 100 hz: 0.006 m/s/s
+# mag:
+# - Total RMS noise: 0.5 milliGauss
+# -------------
+# Movella Dot
+# gyr:
+# - bias error: ?
+# - bias stability: 10 deg/h
+# - noise density: 0.007 deg/s/sqrt(hz)
+# acc:
+# - bias error: ?
+# - bias stability: 30 micro g
+# - noise density: 120 micro g/sqrt(hz)
+# mag:
+# - Total RMS noise: 0.5 milliGauss = 5e-8 Tesla
+
+# units are:
+# - acc: m/s/s
+# - gyr: rad/s
+# - mag: a.u.
+NOISE_LEVELS = {"acc": 0.03, "gyr": jnp.deg2rad(0.5), "mag": 0.01}
+BIAS_LEVELS = {"acc": 0.1, "gyr": jnp.deg2rad(0.5), "mag": 0.0}
+
+
+def add_noise_bias(
+    key: jax.random.PRNGKey,
+    imu_measurements: dict[str, jax.Array],
+    noise_levels: Optional[dict[str, float | None]] = None,
+    bias_levels: Optional[dict[str, float | None]] = None,
+) -> dict[str, jax.Array]:
+    """Add noise and bias to 6D or 9D imu measurements.
+
+    Args:
+        key (jax.random.PRNGKey): Random seed.
+        imu_measurements (dict): IMU measurements without noise and bias.
+            Format is {"gyr": Array, "acc": Array, "mag": Array}.
+
+    Returns:
+        dict: IMU measurements with noise and bias.
+    """
+    noise_levels = {} if noise_levels is None else noise_levels
+    bias_levels = {} if bias_levels is None else bias_levels
+
+    noisy_imu_measurements = {}
+    for sensor in imu_measurements:
+        key, c1, c2 = jax.random.split(key, 3)
+
+        noise_scale = noise_levels.get(sensor, NOISE_LEVELS[sensor])
+        if noise_scale is not None:
+            noise = (
+                jax.random.normal(c1, shape=imu_measurements[sensor].shape)
+                * noise_scale
+            )
+        else:
+            noise = 0.0
+
+        bias_maxval = bias_levels.get(sensor, BIAS_LEVELS[sensor])
+        if bias_maxval is not None:
+            bias = jax.random.uniform(
+                c2, minval=-bias_maxval, maxval=bias_maxval, shape=(3,)
+            )
+        else:
+            bias = 0.0
+
+        noisy_imu_measurements[sensor] = imu_measurements[sensor] + noise + bias
+
+    return noisy_imu_measurements
+
+
+def imu(
+    xs: base.Transform,
+    gravity: jax.Array,
+    dt: float,
+    key: Optional[jax.random.PRNGKey] = None,
+    noisy: bool = False,
+    smoothen_degree: Optional[int] = None,
+    delay: Optional[int] = None,
+    random_s2s_ori: Optional[float] = None,
+    quasi_physical: bool = False,
+    low_pass_filter_pos_f_cutoff: Optional[float] = None,
+    low_pass_filter_rot_cutoff: Optional[float] = None,
+    has_magnetometer: bool = False,
+    magvec: Optional[jax.Array] = None,
+    gyro_second_order: bool = False,
+    natural_units: bool = False,
+    acc_xinyuyi_n: int = 1,
+) -> dict:
+    """Simulates a 6D IMU, `xs` should be Transforms from eps-to-imu.
+    NOTE: `smoothen_degree` is used as window size for moving average.
+    NOTE: If `smoothen_degree` is given, and `delay` is not, then delay is chosen
+    such moving average window is delayed to just be causal.
+    """
+    assert xs.ndim() == 2
+
+    if random_s2s_ori is not None:
+        assert key is not None, "`random_s2s_ori` requires a random seed via `key`"
+        # `xs` are now from eps-to-segment, so add another final rotation from
+        # segment-to-sensor where this transform is only rotational
+        key, consume = jax.random.split(key)
+        xs_s2s = base.Transform.create(
+            rot=maths.quat_random(consume, maxval=random_s2s_ori)
+        )
+        xs = jax.vmap(algebra.transform_mul, in_axes=(None, 0))(xs_s2s, xs)
+
+    if quasi_physical:
+        xs = _quasi_physical_simulation(xs, dt)
+
+    if low_pass_filter_pos_f_cutoff is not None:
+        xs = xs.replace(
+            pos=_butterworth(
+                xs.pos, f_sampling=1 / dt, f_cutoff=low_pass_filter_pos_f_cutoff
+            )
+        )
+
+    if low_pass_filter_rot_cutoff is not None:
+        xs = xs.replace(
+            rot=maths.quat_lowpassfilter(
+                xs.rot, cutoff_freq=low_pass_filter_rot_cutoff, samp_freq=1 / dt
+            )
+        )
+
+    measurements = {
+        "acc": accelerometer(xs, gravity, dt, acc_xinyuyi_n),
+        "gyr": gyroscope(xs.rot, dt, gyro_second_order),
+    }
+
+    if has_magnetometer:
+        if magvec is None:
+            assert key is not None
+            key, consume = jax.random.split(key)
+            magvec = _draw_random_magvec(consume)
+        measurements["mag"] = magnetometer(xs.rot, magvec)
+
+    if smoothen_degree is not None:
+        measurements = jax.tree_map(
+            lambda arr: _moving_average(arr, smoothen_degree),
+            measurements,
+        )
+
+        # if you low-pass filter the imu measurements through a moving average which
+        # effectively uses future values, then it also makes sense to delay the imu
+        # measurements by this amount such that no future information is used
+        if delay is None:
+            half_window = (smoothen_degree - 1) // 2
+            delay = half_window
+
+    if delay is not None and delay > 0:
+        measurements = jax.tree_map(
+            lambda arr: (jnp.pad(arr, ((delay, 0), (0, 0)))[:-delay]), measurements
+        )
+
+    if noisy:
+        assert key is not None, "For noisy sensors random seed `key` must be provided."
+        measurements = add_noise_bias(key, measurements)
+
+    if natural_units:
+        measurements = rescale_natural_units(measurements)
+
+    return measurements
+
+
+_rescale_natural_units_fns = defaultdict(lambda: (lambda arr: arr))
+_rescale_natural_units_fns["gyr"] = lambda gyr: gyr / jnp.pi
+_rescale_natural_units_fns["acc"] = lambda acc: acc / 9.81
+
+
+def rescale_natural_units(X: dict[str, jax.Array]):
+    return {key: _rescale_natural_units_fns[key](val) for key, val in X.items()}
+
+
+def rel_pose(
+    sys: base.System, xs: base.Transform, sys_xs: Optional[base.System] = None
+) -> dict:
+    """Relative pose of the entire system. `sys_scan` defines the parent-child ordering,
+    relative pose is from child to parent in local coordinates. Bodies that connect
+    to the base are skipped (that would be absolute pose).
+
+    Args:
+        sys_scan (base.System): System defining parent-child ordering.
+        xs (base.Transform): Body transforms from base to body.
+        sys_xs (base.System): System that defines the stacking order of `xs`.
+
+    Returns:
+        dict: Child-to-parent quaternions
+    """
+    if sys_xs is None:
+        sys_xs = sys
+
+    if xs.pos.ndim == 3:
+        # swap (n_timesteps, n_links) axes
+        xs = xs.transpose([1, 0, 2])
+
+    assert xs.batch_dim() == sys_xs.num_links()
+
+    qrel = lambda q1, q2: maths.quat_mul(q1, maths.quat_inv(q2))
+
+    y = {}
+
+    def pose_child_to_parent(_, __, name_i: str, p: int):
+        # body connects to base
+        if p == -1:
+            return
+
+        name_p = sys.idx_to_name(p)
+
+        # find the transforms of those named bodies
+        i = sys_xs.name_to_idx(name_i)
+        p = sys_xs.name_to_idx(name_p)
+
+        # get those transforms
+        q1, q2 = xs.take(p).rot, xs.take(i).rot
+
+        y[name_i] = qrel(q1, q2)
+
+    sys.scan(pose_child_to_parent, "ll", sys.link_names, sys.link_parents)
+
+    return y
+
+
+def root_incl(
+    sys: base.System, x: base.Transform, sys_x: base.System
+) -> dict[str, jax.Array]:
+    # (time, nlinks, 4) -> (nlinks, time, 4)
+    rots = x.rot.transpose((1, 0, 2))
+    l_map = sys_x.idx_map("l")
+
+    y = dict()
+
+    def f(_, __, name: str, parent: int):
+        if parent != -1:
+            return
+        y[name] = maths.quat_project(rots[l_map[name]], jnp.array([0.0, 0, 1]))[1]
+
+    sys.scan(f, "ll", sys.link_names, sys.link_parents)
+
+    return y
+
+
+def root_full(
+    sys: base.System, x: base.Transform, sys_x: base.System
+) -> dict[str, jax.Array]:
+    # (time, nlinks, 4) -> (nlinks, time, 4)
+    rots = x.rot.transpose((1, 0, 2))
+    l_map = sys_x.idx_map("l")
+
+    y = dict()
+
+    def f(_, __, name: str, parent: int):
+        if parent != -1:
+            return
+        y[name] = rots[l_map[name]]
+
+    sys.scan(f, "ll", sys.link_names, sys.link_parents)
+
+    return y
+
+
+def joint_axes(
+    sys: base.System,
+    xs: base.Transform,
+    sys_xs: base.System,
+    key: Optional[jax.Array] = None,
+    noisy: bool = False,
+    from_sys: bool = False,
+    randomly_flip: bool = False,
+):
+    """
+    The joint-axes to world is always zeros.
+    """
+    if key is None:
+        assert not noisy
+        assert not randomly_flip
+
+    N = xs.shape(axis=0)
+
+    if from_sys:
+        X = _joint_axes_from_sys(sys, N)
+    else:
+        X = _joint_axes_from_xs(sys, xs, sys_xs)
+
+    if noisy:
+        for name in X:
+            key, c1, c2 = jax.random.split(key, 3)
+            bias = maths.quat_random(c1, maxval=jnp.deg2rad(5.0))
+            noise = maths.quat_random(c2, (N,), maxval=jnp.deg2rad(2.0))
+            dist = maths.quat_mul(noise, bias)
+            X[name]["joint_axes"] = maths.rotate(X[name]["joint_axes"], dist)
+
+    # joint axes to world must be zeros
+    for name, p in zip(sys.link_names, sys.link_parents):
+        if p == -1:
+            X[name]["joint_axes"] = jnp.zeros((N, 3))
+        else:
+            if randomly_flip:
+                key, consume = jax.random.split(key)
+                X[name]["joint_axes"] = (
+                    jax.random.choice(consume, jnp.array([1.0, -1.0]))
+                    * X[name]["joint_axes"]
+                )
+
+    return X
+
+
+def _joint_axes_from_xs(sys, xs, sys_xs):
+
+    xs = sim2real.match_xs(sys, xs, sys_xs)
+
+    _, transform2_rot = sim2real.unzip_xs(sys, xs)
+    qs = transform2_rot.rot.transpose((1, 0, 2))
+
+    l2norm = lambda x: jnp.sqrt(jnp.sum(x**2, axis=-1))
+
+    @jax.vmap
+    def ensure_axis_convention(qs):
+        axis = qs[..., 1:] / (
+            jnp.linalg.norm(qs[..., 1:], axis=-1, keepdims=True) + 1e-6
+        )
+        convention = axis[0]
+        cond = (l2norm(convention - axis) > l2norm(convention + axis))[..., None]
+        return jnp.where(cond, -axis, axis)
+
+    axes = ensure_axis_convention(qs)
+
+    # TODO
+    # not ideal to average vectors that live on a sphere
+    N = axes.shape[1]
+    axes_average = jnp.mean(axes, axis=1)
+    axes_average /= jnp.linalg.norm(axes_average, axis=-1, keepdims=True)
+    axes = jnp.repeat(axes_average[:, None], N, axis=1)
+
+    X = {name: {"joint_axes": axes[sys.name_to_idx(name)]} for name in sys.link_names}
+    return X
+
+
+def _joint_axes_from_sys(sys: base.Transform, N: int) -> dict:
+    "`sys` should be `sys_noimu`. `N` is number of timesteps"
+    xaxis = jnp.array([1.0, 0, 0])
+    yaxis = jnp.array([0.0, 1, 0])
+    zaxis = jnp.array([0.0, 0, 1])
+    id_to_axis = {"x": xaxis, "y": yaxis, "z": zaxis}
+    X = {}
+
+    def f(_, __, name, link_type, link):
+        joint_params = link.joint_params
+        if link_type in ["rx", "ry", "rz"]:
+            joint_axes = id_to_axis[link_type[1]]
+        elif link_type == "rr":
+            joint_axes = joint_params["rr"]["joint_axes"]
+        elif link_type[:6] == "rr_imp":
+            joint_axes = joint_params[link_type]["joint_axes"]
+        else:
+            joint_axes = xaxis
+        X[name] = {"joint_axes": joint_axes}
+
+    sys.scan(f, "lll", sys.link_names, sys.link_types, sys.links)
+    X = jax.tree_map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
+    return X
+
+
+def _moving_average(arr: jax.Array, window: int) -> jax.Array:
+    "Padds with left and right values of array."
+    assert window % 2 == 1
+    assert window > 1, "Window size of 1 would be a no-op"
+    arr_smooth = jnp.zeros((len(arr) + window - 1,) + arr.shape[1:])
+    half_window = (window - 1) // 2
+    arr_padded = arr_smooth.at[half_window : (len(arr) + half_window)].set(arr)
+    arr_padded = arr_padded.at[:half_window].set(arr[0])
+    arr_padded = arr_padded.at[-half_window:].set(arr[-1])
+
+    for i in range(-half_window, half_window + 1):
+        rolled = jnp.roll(arr_padded, i, axis=0)
+        arr_smooth += rolled
+    arr_smooth = arr_smooth / window
+    return arr_smooth[half_window : (len(arr) + half_window)]
+
+
+_quasi_physical_sys_str = r"""
+<x_xy>
+    <options gravity="0 0 0"/>
+    <worldbody>
+        <body name="IMU" joint="p3d" damping="0.1 0.1 0.1" spring_stiff="3 3 3">
+            <geom type="box" mass="0.002" dim="0.01 0.01 0.01"/>
+        </body>
+    </worldbody>
+</x_xy>
+"""
+
+
+def _quasi_physical_simulation_beautiful(
+    xs: base.Transform, dt: float
+) -> base.Transform:
+    sys = io.load_sys_from_str(_quasi_physical_sys_str).replace(dt=dt)
+
+    def step_dynamics(state: base.State, x):
+        state = algorithms.step(sys.replace(link_spring_zeropoint=x.pos), state)
+        return state, state.q
+
+    state = base.State.create(sys, q=xs.pos[0])
+    _, pos = jax.lax.scan(step_dynamics, state, xs)
+    return xs.replace(pos=pos)
+
+
+_constants = {
+    "qp_damp": 35.0,
+    "qp_stif": 625.0,
+}
+
+
+def _quasi_physical_simulation(xs: base.Transform, dt: float) -> base.Transform:
+    mass = 1.0
+    damp = _constants["qp_damp"]
+    stif = _constants["qp_stif"]
+
+    def step_dynamics(state, zeropoint):
+        pos, vel = state
+        zeropoint_pos, zeropoint_vel = zeropoint
+        acc = (damp * (zeropoint_vel - vel) + stif * (zeropoint_pos - pos)) / mass
+        vel += dt * acc
+        # semi-implicit, so use already next velocity
+        pos += dt * vel
+        return (pos, vel), pos
+
+    zero_vel = jnp.zeros_like(xs.pos[0])
+    state = (xs.pos[0], zero_vel)
+    zeropoint_vel = jnp.vstack((zero_vel, jnp.diff(xs.pos, axis=0) / dt))
+    zeropoint = (xs.pos, zeropoint_vel)
+    _, pos = jax.lax.scan(step_dynamics, state, zeropoint)
+    return xs.replace(pos=pos)
+
+
+def _butterworth(
+    signal: jax.Array,
+    f_sampling: float,
+    f_cutoff: int,
+    method: str = "forward_backward",
+) -> jax.Array:
+    """https://stackoverflow.com/questions/20924868/calculate-coefficients-of-2nd-order
+    -butterworth-low-pass-filter"""
+
+    if method == "forward_backward":
+        signal = _butterworth(signal, f_sampling, f_cutoff, "forward")
+        return _butterworth(signal, f_sampling, f_cutoff, "backward")
+    elif method == "forward":
+        pass
+    elif method == "backward":
+        signal = jnp.flip(signal, axis=0)
+    else:
+        raise NotImplementedError
+
+    ff = f_cutoff / f_sampling
+    ita = 1.0 / jnp.tan(jnp.pi * ff)
+    q = jnp.sqrt(2.0)
+    b0 = 1.0 / (1.0 + q * ita + ita**2)
+    b1 = 2 * b0
+    b2 = b0
+    a1 = 2.0 * (ita**2 - 1.0) * b0
+    a2 = -(1.0 - q * ita + ita**2) * b0
+
+    def f(carry, x_i):
+        x_im1, x_im2, y_im1, y_im2 = carry
+        y_i = b0 * x_i + b1 * x_im1 + b2 * x_im2 + a1 * y_im1 + a2 * y_im2
+        return (x_i, x_im1, y_i, y_im1), y_i
+
+    init = (signal[1], signal[0]) * 2
+    signal = jax.lax.scan(f, init, signal[2:])[1]
+    signal = jnp.concatenate((signal[0:1],) * 2 + (signal,))
+
+    if method == "backward":
+        signal = jnp.flip(signal, axis=0)
+
+    return signal