PyPI - imt-ring - Versions diffs - 1.2.1__py3-none-any.whl - Mend

imt-ring 1.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (83) hide show

imt_ring-1.2.1.dist-info/METADATA +91 -0
imt_ring-1.2.1.dist-info/RECORD +83 -0
imt_ring-1.2.1.dist-info/WHEEL +5 -0
imt_ring-1.2.1.dist-info/top_level.txt +1 -0
ring/__init__.py +63 -0
ring/algebra.py +100 -0
ring/algorithms/__init__.py +45 -0
ring/algorithms/_random.py +403 -0
ring/algorithms/custom_joints/__init__.py +6 -0
ring/algorithms/custom_joints/rr_imp_joint.py +69 -0
ring/algorithms/custom_joints/rr_joint.py +33 -0
ring/algorithms/custom_joints/suntay.py +424 -0
ring/algorithms/dynamics.py +345 -0
ring/algorithms/generator/__init__.py +25 -0
ring/algorithms/generator/base.py +414 -0
ring/algorithms/generator/batch.py +282 -0
ring/algorithms/generator/motion_artifacts.py +222 -0
ring/algorithms/generator/pd_control.py +182 -0
ring/algorithms/generator/randomize.py +119 -0
ring/algorithms/generator/transforms.py +410 -0
ring/algorithms/generator/types.py +36 -0
ring/algorithms/jcalc.py +840 -0
ring/algorithms/kinematics.py +202 -0
ring/algorithms/sensors.py +582 -0
ring/base.py +1046 -0
ring/io/__init__.py +9 -0
ring/io/examples/branched.xml +24 -0
ring/io/examples/exclude/knee_trans_dof.xml +26 -0
ring/io/examples/exclude/standard_sys.xml +106 -0
ring/io/examples/exclude/standard_sys_rr_imp.xml +106 -0
ring/io/examples/inv_pendulum.xml +14 -0
ring/io/examples/knee_flexible_imus.xml +22 -0
ring/io/examples/spherical_stiff.xml +11 -0
ring/io/examples/symmetric.xml +12 -0
ring/io/examples/test_all_1.xml +39 -0
ring/io/examples/test_all_2.xml +39 -0
ring/io/examples/test_ang0_pos0.xml +9 -0
ring/io/examples/test_control.xml +16 -0
ring/io/examples/test_double_pendulum.xml +14 -0
ring/io/examples/test_free.xml +11 -0
ring/io/examples/test_kinematics.xml +23 -0
ring/io/examples/test_morph_system/four_seg_seg1.xml +26 -0
ring/io/examples/test_morph_system/four_seg_seg3.xml +26 -0
ring/io/examples/test_randomize_position.xml +26 -0
ring/io/examples/test_sensors.xml +13 -0
ring/io/examples/test_three_seg_seg2.xml +23 -0
ring/io/examples.py +42 -0
ring/io/test_examples.py +6 -0
ring/io/xml/__init__.py +6 -0
ring/io/xml/abstract.py +300 -0
ring/io/xml/from_xml.py +299 -0
ring/io/xml/test_from_xml.py +56 -0
ring/io/xml/test_to_xml.py +31 -0
ring/io/xml/to_xml.py +94 -0
ring/maths.py +397 -0
ring/ml/__init__.py +33 -0
ring/ml/base.py +292 -0
ring/ml/callbacks.py +434 -0
ring/ml/ml_utils.py +272 -0
ring/ml/optimizer.py +149 -0
ring/ml/params/0x13e3518065c21cd8.pickle +0 -0
ring/ml/ringnet.py +279 -0
ring/ml/train.py +318 -0
ring/ml/training_loop.py +131 -0
ring/rendering/__init__.py +2 -0
ring/rendering/base_render.py +271 -0
ring/rendering/mujoco_render.py +222 -0
ring/rendering/vispy_render.py +340 -0
ring/rendering/vispy_visuals.py +290 -0
ring/sim2real/__init__.py +7 -0
ring/sim2real/sim2real.py +288 -0
ring/spatial.py +126 -0
ring/sys_composer/__init__.py +5 -0
ring/sys_composer/delete_sys.py +114 -0
ring/sys_composer/inject_sys.py +110 -0
ring/sys_composer/morph_sys.py +361 -0
ring/utils/__init__.py +21 -0
ring/utils/batchsize.py +51 -0
ring/utils/colab.py +48 -0
ring/utils/hdf5.py +198 -0
ring/utils/normalizer.py +56 -0
ring/utils/path.py +44 -0
ring/utils/utils.py +161 -0

ring/io/xml/test_from_xml.py ADDED Viewed

@@ -0,0 +1,56 @@
+import jax.numpy as jnp
+import ring
+sys_str = """
+<x_xy model="model">
+<options gravity=".1 2 3" dt=".03"/>
+    <worldbody>
+        <body name="name" joint="rx" pos="1 2 3" euler="30 30 30" damping=".7" armature=".8" spring_stiff="1" spring_zero=".9">
+            <geom type="box" mass="2.7" dim="0.2 0.3 0.4" color="black" edge_color="pink"/>
+        </body>
+    </worldbody>
+</x_xy>
+"""  # noqa: E501
+def test_from_xml():
+    pos = jnp.array([1.0, 2, 3])
+    sys1 = ring.System(
+        [-1],
+        ring.base.Link(
+            ring.base.Transform(
+                pos=pos,
+                rot=ring.maths.quat_euler(
+                    jnp.array([jnp.deg2rad(30), jnp.deg2rad(30), jnp.deg2rad(30)])
+                ),
+            ),
+            pos_min=pos,
+            pos_max=pos,
+        ).batch(),
+        ["rx"],
+        link_damping=jnp.array([0.7]),
+        link_armature=jnp.array([0.8]),
+        link_spring_zeropoint=jnp.array([0.9]),
+        link_spring_stiffness=jnp.array([1.0]),
+        dt=0.03,
+        geoms=[
+            ring.base.Box(
+                jnp.array(2.7),
+                ring.Transform.zero(),
+                0,
+                "black",
+                "pink",
+                jnp.array(0.2),
+                jnp.array(0.3),
+                jnp.array(0.4),
+            )
+        ],
+        gravity=jnp.array([0.1, 2, 3.0]),
+        link_names=["name"],
+        model_name="model",
+        omc=[None],
+    )
+    sys1 = sys1.parse()
+    sys2 = ring.io.load_sys_from_str(sys_str)
+    assert ring.utils.sys_compare(sys1, sys2)

ring/io/xml/test_to_xml.py ADDED Viewed

@@ -0,0 +1,31 @@
+import logging
+import ring
+from ring.base import System
+from ring.utils import sys_compare
+def test_save_sys_to_str():
+    for original_sys in ring.io.list_load_examples():
+        sys_to_xml_str = ring.io.save_sys_to_str(original_sys)
+        logging.debug(sys_to_xml_str)
+        compare_sys = ring.io.load_sys_from_str(sys_to_xml_str)
+        assert sys_compare(
+            original_sys, compare_sys
+        ), f"Failed {original_sys.model_name}.xml"
+        print(f"Passed {original_sys.model_name}.xml")
+    def double_load_xml_to_sys(example: str) -> System:
+        orig_sys = ring.io.load_example(example)
+        exported_xml = ring.io.save_sys_to_str(orig_sys)
+        new_sys = ring.io.load_sys_from_str(exported_xml)
+        return new_sys
+    sys_test_xml_1 = double_load_xml_to_sys("test_all_1.xml")
+    sys_test_xml_2 = double_load_xml_to_sys("test_all_2.xml")
+    assert not sys_compare(sys_test_xml_1, sys_test_xml_2)

ring/io/xml/to_xml.py ADDED Viewed

@@ -0,0 +1,94 @@
+import warnings
+from xml.dom.minidom import parseString
+from xml.etree.ElementTree import Element
+from xml.etree.ElementTree import SubElement
+from xml.etree.ElementTree import tostring
+import jax.numpy as jnp
+from ring import base
+from tree_utils import batch_concat
+from . import abstract
+from .abstract import _to_str
+def save_sys_to_str(sys: base.System) -> str:
+    for joint_type in sys.links.joint_params:
+        for i, link_name in enumerate(sys.link_names):
+            joint_params_flat = batch_concat((sys.links[i]).joint_params[joint_type], 0)
+            if not jnp.all(joint_params_flat == 0.0):
+                warnings.warn(
+                    "The system has `sys.links.joint_params` unequal to the 'default'"
+                    f" value (of zeros). In particular the link `{link_name}` has for"
+                    f" the jointtype `{joint_type}` the values {joint_params_flat}. "
+                    "This will not be preserved in the xml."
+                )
+    global_index_map = {qd: sys.idx_map(qd) for qd in ["q", "d"]}
+    # Create root element
+    x_xy = Element("x_xy")
+    x_xy.set("model", sys.model_name)
+    options = SubElement(x_xy, "options")
+    options.set("dt", str(sys.dt))
+    options.set("gravity", _to_str(sys.gravity))
+    # Create worldbody
+    worldbody = SubElement(x_xy, "worldbody")
+    def process_link(link_idx: int, parent_elem: Element):
+        link = sys.links[link_idx]
+        link_typ = sys.link_types[link_idx]
+        link_name = sys.link_names[link_idx]
+        # Create body element
+        body = SubElement(parent_elem, "body")
+        body.set("joint", link_typ)
+        body.set("name", link_name)
+        # Set attributes
+        abstract.AbsTrans.to_xml(body, link.transform1)
+        abstract.AbsPosMinMax.to_xml(body, link.pos_min, link.pos_max)
+        abstract.AbsDampArmaStiffZero.to_xml(
+            body,
+            sys.link_damping[global_index_map["d"][link_name]],
+            sys.link_armature[global_index_map["d"][link_name]],
+            sys.link_spring_stiffness[global_index_map["d"][link_name]],
+            sys.link_spring_zeropoint[global_index_map["q"][link_name]],
+            base.Q_WIDTHS[link_typ],
+            base.QD_WIDTHS[link_typ],
+            link_typ,
+        )
+        # Add geometry elements
+        geoms = sys.geoms
+        for geom in geoms:
+            if geom.link_idx == link_idx:
+                geom_elem = SubElement(body, "geom")
+                abstract_class = abstract.geometry_to_abstract[type(geom)]
+                abstract_class.to_xml(geom_elem, geom)
+        # Maybe add omc element
+        omc_link = sys.omc[link_idx]
+        if omc_link is not None:
+            omc_elem = SubElement(body, "omc")
+            abstract.AbsMaxCoordOMC.to_xml(omc_elem, omc_link)
+        # Recursively process child links
+        for child_idx, parent_idx in enumerate(sys.link_parents):
+            if parent_idx == link_idx:
+                process_link(child_idx, body)
+    for root_link_idx, parent_idx in enumerate(sys.link_parents):
+        if parent_idx == -1:
+            process_link(root_link_idx, worldbody)
+    # Pretty print xml
+    xml_str = parseString(tostring(x_xy)).toprettyxml(indent="  ")
+    return xml_str
+def save_sys_to_xml(sys: base.System, xml_path: str) -> None:
+    xml_str = save_sys_to_str(sys)
+    with open(xml_path, "w") as f:
+        f.write(xml_str)

ring/maths.py ADDED Viewed

@@ -0,0 +1,397 @@
+from functools import partial
+import jax
+from jax import custom_jvp
+import jax.numpy as jnp
+import jax.random as jrand
+def wrap_to_pi(phi):
+    "Wraps angle `phi` (radians) to interval [-pi, pi]."
+    return (phi + jnp.pi) % (2 * jnp.pi) - jnp.pi
+x_unit_vector = jnp.array([1.0, 0, 0])
+y_unit_vector = jnp.array([0.0, 1, 0])
+z_unit_vector = jnp.array([0.0, 0, 1])
+def unit_vectors(xyz: int | str):
+    if isinstance(xyz, str):
+        xyz = {"x": 0, "y": 1, "z": 2}[xyz]
+    return [x_unit_vector, y_unit_vector, z_unit_vector][xyz]
+@partial(jnp.vectorize, signature="(k)->(1)")
+def safe_norm(x):
+    """Grad-safe for x=0.0. Norm along last axis."""
+    assert x.ndim == 1
+    is_zero = jnp.all(jnp.isclose(x, 0.0), axis=-1, keepdims=False)
+    return jax.lax.cond(
+        is_zero,
+        lambda x: jnp.array([0.0], dtype=x.dtype),
+        lambda x: jnp.linalg.norm(x, keepdims=True),
+        x,
+    )
+@partial(jnp.vectorize, signature="(k)->(k)")
+def safe_normalize(x):
+    """Execution- and Grad-safe for x=0.0. Normalizes along last axis."""
+    assert x.ndim == 1
+    is_zero = jnp.allclose(x, 0.0)
+    return jax.lax.cond(
+        is_zero,
+        lambda x: jnp.zeros_like(x),
+        lambda x: x / jnp.where(is_zero, 1.0, safe_norm(x)),
+        x,
+    )
+@custom_jvp
+def safe_arccos(x: jnp.ndarray) -> jnp.ndarray:
+    """Trigonometric inverse cosine, element-wise with safety clipping in grad."""
+    return jnp.arccos(x)
+@safe_arccos.defjvp
+def _safe_arccos_jvp(primal, tangent):
+    (x,) = primal
+    (x_dot,) = tangent
+    primal_out = safe_arccos(x)
+    tangent_out = -x_dot / jnp.sqrt(1.0 - jnp.clip(x, -1 + 1e-7, 1 - 1e-7) ** 2.0)
+    return primal_out, tangent_out
+@custom_jvp
+def safe_arcsin(x: jnp.ndarray) -> jnp.ndarray:
+    """Trigonometric inverse sine, element-wise with safety clipping in grad."""
+    return jnp.arcsin(x)
+@safe_arcsin.defjvp
+def _safe_arcsin_jvp(primal, tangent):
+    (x,) = primal
+    (x_dot,) = tangent
+    primal_out = safe_arccos(x)
+    tangent_out = x_dot / jnp.sqrt(1.0 - jnp.clip(x, -1 + 1e-7, 1 - 1e-7) ** 2.0)
+    return primal_out, tangent_out
+@partial(jnp.vectorize, signature="(4)->(4)")
+def ensure_positive_w(q):
+    return jnp.where(q[0] < 0, -q, q)
+def angle_error(q, qhat):
+    "Absolute angle in radians between `q` and `qhat`."
+    return jnp.abs(quat_angle(quat_mul(quat_inv(q), qhat)))
+def unit_quats_like(array):
+    "Array of *unit* quaternions of identical shape."
+    if array.shape[-1] != 4:
+        raise Exception()
+    return jnp.ones(array.shape[:-1])[..., None] * jnp.array([1.0, 0, 0, 0])
+@partial(jnp.vectorize, signature="(4),(4)->(4)")
+def quat_mul(u: jnp.ndarray, v: jnp.ndarray) -> jnp.ndarray:
+    "Multiplies two quaternions."
+    q = jnp.array(
+        [
+            u[0] * v[0] - u[1] * v[1] - u[2] * v[2] - u[3] * v[3],
+            u[0] * v[1] + u[1] * v[0] + u[2] * v[3] - u[3] * v[2],
+            u[0] * v[2] - u[1] * v[3] + u[2] * v[0] + u[3] * v[1],
+            u[0] * v[3] + u[1] * v[2] - u[2] * v[1] + u[3] * v[0],
+        ]
+    )
+    return q
+def quat_inv(q: jnp.ndarray) -> jnp.ndarray:
+    "Calculates the inverse of quaternion q."
+    return q * jnp.array([1.0, -1, -1, -1])
+@partial(jnp.vectorize, signature="(3),(4)->(3)")
+def rotate(vector: jnp.ndarray, quat: jnp.ndarray):
+    """Rotates a vector `vector` by a *unit* quaternion `quat`."""
+    qvec = jnp.array([0, *vector])
+    return rotate_quat(qvec, quat)[1:4]
+def rotate_matrix(matrix: jax.Array, quat: jax.Array):
+    "Rotate matrix `matrix` by a *unit* quaternion `quat`."
+    E = quat_to_3x3(quat)
+    return E @ matrix @ E.T
+def rotate_quat(q: jax.Array, quat: jax.Array):
+    "Rotate quaternion `q` by `quat`"
+    return quat_mul(quat, quat_mul(q, quat_inv(quat)))
+@partial(jnp.vectorize, signature="(3),()->(4)")
+def quat_rot_axis(axis: jnp.ndarray, angle: jnp.ndarray) -> jnp.ndarray:
+    """Construct a *unit* quaternion that describes rotating around
+    `axis` by `angle` (radians).
+    This is the interpretation of rotating the vector and *not*
+    the frame.
+    For the interpretation of rotating the frame and *not* the
+    vector, you should use angle -> -angle.
+    NOTE: Usually, we actually want the second interpretation. Think about it,
+    we use quaternions to re-express vectors in other frames. But the
+    vectors stay the same. We only transform them to a common frames.
+    """
+    assert axis.shape == (3,)
+    assert angle.shape == ()
+    axis = safe_normalize(axis)
+    # NOTE: CONVENTION
+    # 23.04.23
+    # this fixes the issue of prismatic joints being inverted w.r.t.
+    # gravity vector.
+    # The reason is that it inverts the way how revolute joints behave
+    # Such that prismatic joints work by inverting gravity
+    angle *= -1.0
+    s, c = jnp.sin(angle / 2), jnp.cos(angle / 2)
+    return jnp.array([c, *(axis * s)])
+@partial(jnp.vectorize, signature="(3,3)->(4)")
+def quat_from_3x3(m: jnp.ndarray) -> jnp.ndarray:
+    """Converts 3x3 rotation matrix to *unit* quaternion."""
+    w = jnp.sqrt(1 + m[0, 0] + m[1, 1] + m[2, 2]) / 2.0
+    x = (m[2][1] - m[1][2]) / (w * 4)
+    y = (m[0][2] - m[2][0]) / (w * 4)
+    z = (m[1][0] - m[0][1]) / (w * 4)
+    return jnp.array([w, x, y, z])
+@partial(jnp.vectorize, signature="(4)->(3,3)")
+def quat_to_3x3(q: jnp.ndarray) -> jnp.ndarray:
+    """Converts *unit* quaternion to 3x3 rotation matrix."""
+    d = jnp.dot(q, q)
+    w, x, y, z = q
+    s = 2 / d
+    xs, ys, zs = x * s, y * s, z * s
+    wx, wy, wz = w * xs, w * ys, w * zs
+    xx, xy, xz = x * xs, x * ys, x * zs
+    yy, yz, zz = y * ys, y * zs, z * zs
+    return jnp.array(
+        [
+            jnp.array([1 - (yy + zz), xy - wz, xz + wy]),
+            jnp.array([xy + wz, 1 - (xx + zz), yz - wx]),
+            jnp.array([xz - wy, yz + wx, 1 - (xx + yy)]),
+        ]
+    )
+def quat_random(
+    key: jrand.PRNGKey, batch_shape: tuple = (), maxval: float = jnp.pi
+) -> jax.Array:
+    """Provides a random *unit* quaternion, sampled uniformly"""
+    assert key.shape == (2,), f"{key.shape}"
+    shape = batch_shape + (4,)
+    qs = safe_normalize(jrand.normal(key, shape))
+    def _scale_angle():
+        axis, angle = quat_to_rot_axis(qs)
+        angle_scaled = angle * maxval / jnp.pi
+        return quat_rot_axis(axis, angle_scaled)
+    return jax.lax.cond(maxval == jnp.pi, lambda: qs, _scale_angle)
+def quat_euler(angles, intrinsic=True, convention="zyx"):
+    "Construct a *unit* quaternion from Euler angles (radians)."
+    @partial(jnp.vectorize, signature="(3)->(4)")
+    def _quat_euler(angles):
+        xunit = jnp.array([1.0, 0.0, 0.0])
+        yunit = jnp.array([0.0, 1.0, 0.0])
+        zunit = jnp.array([0.0, 0.0, 1.0])
+        axes_map = {
+            "x": xunit,
+            "y": yunit,
+            "z": zunit,
+        }
+        q1 = quat_rot_axis(axes_map[convention[0]], angles[0])
+        q2 = quat_rot_axis(axes_map[convention[1]], angles[1])
+        q3 = quat_rot_axis(axes_map[convention[2]], angles[2])
+        if intrinsic:
+            return quat_mul(q3, quat_mul(q2, q1))
+        else:
+            return quat_mul(q1, quat_mul(q2, q3))
+    return _quat_euler(angles)
+@partial(jnp.vectorize, signature="(4)->()")
+def quat_angle(q):
+    "Extract rotation angle (radians) of quaternion `q`."
+    phi = 2 * jnp.arctan2(safe_norm(q[1:])[0], q[0])
+    return wrap_to_pi(phi)
+def quat_angle_constantAxisOverTime(qs):
+    assert qs.ndim == 2
+    assert qs.shape[-1] == 4
+    l2norm = lambda x: jnp.sqrt(jnp.sum(x**2, axis=-1))
+    axis = safe_normalize(qs[:, 1:])
+    angle = quat_angle(qs)[:, None]
+    convention = axis[0]
+    cond = (l2norm(convention - axis) > l2norm(convention + axis))[..., None]
+    return jnp.where(cond, -angle, angle)[:, 0]
+@partial(jnp.vectorize, signature="(4)->(3),()")
+def quat_to_rot_axis(q):
+    "Extract unit-axis and angle from quaternion `q`."
+    angle = quat_angle(q)
+    # NOTE: CONVENTION
+    angle *= -1.0
+    axis = safe_normalize(q[1:])
+    return axis, angle
+@partial(jnp.vectorize, signature="(3)->(4)")
+def euler_to_quat(angles: jnp.ndarray) -> jnp.ndarray:
+    """Converts euler rotations in radians to quaternion."""
+    # this follows the Tait-Bryan intrinsic rotation formalism: x-y'-z''
+    c1, c2, c3 = jnp.cos(angles / 2)
+    s1, s2, s3 = jnp.sin(angles / 2)
+    w = c1 * c2 * c3 - s1 * s2 * s3
+    x = s1 * c2 * c3 + c1 * s2 * s3
+    y = c1 * s2 * c3 - s1 * c2 * s3
+    z = c1 * c2 * s3 + s1 * s2 * c3
+    # NOTE: CONVENTION
+    return quat_inv(jnp.array([w, x, y, z]))
+@partial(jnp.vectorize, signature="(4)->(3)")
+def quat_to_euler(q: jnp.ndarray) -> jnp.ndarray:
+    """Converts quaternions to euler rotations in radians."""
+    # this follows the Tait-Bryan intrinsic rotation formalism: x-y'-z''
+    # NOTE: CONVENTION
+    q = quat_inv(q)
+    z = jnp.arctan2(
+        -2 * q[1] * q[2] + 2 * q[0] * q[3],
+        q[1] * q[1] + q[0] * q[0] - q[3] * q[3] - q[2] * q[2],
+    )
+    # TODO: Investigate why quaternions go so big we need to clip.
+    y = safe_arcsin(jnp.clip(2 * q[1] * q[3] + 2 * q[0] * q[2], -1.0, 1.0))
+    x = jnp.arctan2(
+        -2 * q[2] * q[3] + 2 * q[0] * q[1],
+        q[3] * q[3] - q[2] * q[2] - q[1] * q[1] + q[0] * q[0],
+    )
+    return jnp.array([x, y, z])
+@partial(jnp.vectorize, signature="(4),(3)->(4),(4)")
+def quat_project(q: jax.Array, k: jax.Array) -> tuple[jax.Array, jax.Array]:
+    """Decompose quaternion into a primary rotation around axis `k` such that
+    the residual rotation's angle is minimized.
+    Args:
+        q (jax.Array): Quaternion to decompose.
+        k (jax.Array): Primary axis direction.
+    Returns:
+        tuple[jax.Array, jax.Array]: Primary quaternion, residual quaternion
+    """
+    phi_pri = 2 * jnp.arctan2(q[1:] @ k, q[0])
+    # NOTE: CONVENTION
+    q_pri = quat_rot_axis(k, -phi_pri)
+    q_res = quat_mul(q, quat_inv(q_pri))
+    return q_pri, q_res
+def quat_avg(qs: jax.Array):
+    "Tolga Birdal's algorithm."
+    if qs.ndim == 1:
+        qs = qs[None, :]
+    assert qs.ndim == 2
+    return jnp.linalg.eigh(
+        jnp.einsum("ij,ik,i->...jk", qs, qs, jnp.ones((qs.shape[0],)))
+    )[1][:, -1]
+# cutoff_freq=20.0; sampe_freq=100.0
+# -> alpha = 0.55686
+# cutoff_freq=15.0
+# -> alpha = 0.48519
+def quat_lowpassfilter(
+    qs: jax.Array,
+    cutoff_freq: float = 20.0,
+    samp_freq: float = 100.0,
+    filtfilt: bool = False,
+) -> jax.Array:
+    assert qs.ndim == 2
+    assert qs.shape[1] == 4
+    if filtfilt:
+        qs = quat_lowpassfilter(qs, cutoff_freq, samp_freq, filtfilt=False)
+        qs = quat_lowpassfilter(jnp.flip(qs, 0), cutoff_freq, samp_freq, filtfilt=False)
+        return jnp.flip(qs, 0)
+    omega_times_Ts = 2 * jnp.pi * cutoff_freq / samp_freq
+    alpha = omega_times_Ts / (1 + omega_times_Ts)
+    def f(y, x):
+        # error quaternion; current state -> target
+        q_err = quat_mul(x, quat_inv(y))
+        # scale down error quaternion
+        axis, angle = quat_to_rot_axis(q_err)
+        # ensure angle >= 0
+        axis, angle = jax.lax.cond(
+            angle < 0,
+            lambda axis, angle: (-axis, -angle),
+            lambda axis, angle: (axis, angle),
+            axis,
+            angle,
+        )
+        angle_scaled = angle * alpha
+        q_err_scaled = quat_rot_axis(axis, angle_scaled)
+        # move small step toward error quaternion
+        y = quat_mul(q_err_scaled, y)
+        return y, y
+    qs_filtered = jax.lax.scan(f, qs[0], qs[1:])[1]
+    # padd with first value, such that length remains equal
+    qs_filtered = jnp.vstack((qs[0:1], qs_filtered))
+    # renormalize due to float32 numerical errors accumulating
+    return qs_filtered / jnp.linalg.norm(qs_filtered, axis=-1, keepdims=True)
+def quat_inclinationAngle(q: jax.Array):
+    head, incl = quat_project(q, jnp.array([0.0, 0, 1]))
+    return quat_angle(incl)
+def quat_headingAngle(q: jax.Array):
+    head, incl = quat_project(q, jnp.array([0.0, 0, 1]))
+    return quat_angle(head)
+def quat_transfer_heading(q_from: jax.Array, q_to: jax.Array):
+    heading = quat_project(q_from, jnp.array([0.0, 0, 1]))[0]
+    # set heading to zero in the `q_to` quaternions
+    q_to = quat_project(q_to, jnp.array([0.0, 0, 1]))[1]
+    return quat_mul(q_to, heading)

ring/ml/__init__.py ADDED Viewed

@@ -0,0 +1,33 @@
+from . import base
+from . import callbacks
+from . import ml_utils
+from . import optimizer
+from . import ringnet
+from . import train
+from . import training_loop
+from .base import AbstractFilter
+from .ml_utils import on_cluster
+from .ml_utils import unique_id
+from .optimizer import make_optimizer
+from .ringnet import RING
+from .train import train_fn
+def RING_ICML24(params=None, **kwargs):
+    """Create the RING network used in the icml24 paper.
+    X[..., :3]  = acc
+    X[..., 3:6] = gyr
+    X[..., 6:9] = jointaxis
+    X[..., 9:]  = dt
+    """
+    from pathlib import Path
+    if params is None:
+        params = Path(__file__).parent.joinpath("params/0x13e3518065c21cd8.pickle")
+    ringnet = RING(params=params, **kwargs)  # noqa: F811
+    ringnet = base.ScaleX_FilterWrapper(ringnet)
+    ringnet = base.LPF_FilterWrapper(ringnet, 10.0, samp_freq=None)
+    ringnet = base.GroundTruthHeading_FilterWrapper(ringnet)
+    return ringnet