imt-ring 1.3.13__tar.gz → 1.4.1__tar.gz

{imt_ring-1.3.13 → imt_ring-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.13
+Version: 1.4.1
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.3.13 → imt_ring-1.4.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.3.13"
+version = "1.4.1"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.3.13 → imt_ring-1.4.1}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.13
+Version: 1.4.1
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.3.13 → imt_ring-1.4.1}/src/imt_ring.egg-info/SOURCES.txt

@@ -68,6 +68,7 @@ src/ring/ml/rnno_v1.py
 src/ring/ml/train.py
 src/ring/ml/training_loop.py
 src/ring/ml/params/0x13e3518065c21cd8.pickle
+src/ring/ml/params/0x1d76628065a71e0f.pickle
 src/ring/rendering/__init__.py
 src/ring/rendering/base_render.py
 src/ring/rendering/mujoco_render.py
@@ -99,6 +100,7 @@ tests/test_maths.py
 tests/test_ml_utils.py
 tests/test_motion_artifacts.py
 tests/test_pd_control.py
+tests/test_quickstart_example.py
 tests/test_random.py
 tests/test_randomize.py
 tests/test_rcmg.py

imt_ring-1.4.1/src/ring/__init__.py

@@ -0,0 +1,143 @@
+from . import algebra
+from . import algorithms
+from . import base
+from . import io
+from . import maths
+from . import ml
+from . import rendering
+from . import sim2real
+from . import spatial
+from . import sys_composer
+from . import utils
+from .algorithms import join_motionconfigs
+from .algorithms import JointModel
+from .algorithms import MotionConfig
+from .algorithms import RCMG
+from .algorithms import register_new_joint_type
+from .algorithms import step
+from .base import State
+from .base import System
+from .base import Transform
+
+
+def RING(lam: list[int], Ts: float | None):
+    """Creates the RING network.
+
+    Params:
+        lam: parent array
+        Ts : sampling interval of IMU data; time delta in seconds
+
+    Usage:
+    >>> import ring
+    >>> import numpy as np
+    >>>
+    >>> T  : int       = 30        # sequence length     [s]
+    >>> Ts : float     = 0.01      # sampling interval   [s]
+    >>> B  : int       = 1         # batch size
+    >>> lam: list[int] = [0, 1, 2] # parent array
+    >>> N  : int       = len(lam)  # number of bodies
+    >>> T_i: int       = int(T/Ts) # number of timesteps
+    >>>
+    >>> X = np.zeros((B, T_i, N, 9))
+    >>> # where X is structured as follows:
+    >>> # X[..., :3]  = acc
+    >>> # X[..., 3:6] = gyr
+    >>> # X[..., 6:9] = jointaxis
+    >>>
+    >>> # let's assume we have an IMU on each outer segment of the
+    >>> # three-segment kinematic chain
+    >>> X[:, :, 0, :3]  = acc_segment1
+    >>> X[:, :, 2, :3]  = acc_segment3
+    >>> X[:, :, 0, 3:6] = gyr_segment1
+    >>> X[:, :, 2, 3:6] = gyr_segment3
+    >>>
+    >>> ringnet = ring.RING(lam, Ts)
+    >>>
+    >>> yhat, _ = ringnet.apply(X)
+    >>> # yhat : unit quaternions, shape = (B, T_i, N, 4)
+    >>>
+    >>> # use `jax.jit` to compile the forward pass
+    >>> jit_apply = jax.jit(ringnet.apply)
+    >>> yhat, _ = jit_apply(X)
+    >>>
+    >>> # manually pass in and out the hidden state like so
+    >>> initial_state = None
+    >>> yhat, state = ringnet.apply(X, state=initial_state)
+    >>> # state: final hidden state, shape = (B, N, 2*H)
+
+    """
+    from pathlib import Path
+    import warnings
+
+    if Ts is not None and (Ts > (1 / 40) or Ts < (1 / 200)):
+        warnings.warn(
+            "RING was only trained on sampling rates between 40 to 200 Hz "
+            f"but found {1 / Ts}Hz"
+        )
+
+    if Ts is not None and Ts == 0.01:
+        # this set of parameters was trained exclusively on 100Hz data; it also
+        # expects F=9 features per node and not F=10 where the last features is
+        # the sampling interval Ts
+        params = Path(__file__).parent.joinpath("ml/params/0x1d76628065a71e0f.pickle")
+        add_Ts = False
+    else:
+        # this set of parameters was trained on sampling rates from 40 to 200 Hz
+        params = Path(__file__).parent.joinpath("ml/params/0x13e3518065c21cd8.pickle")
+        add_Ts = True
+
+    ringnet = ml.RING(params=params, lam=tuple(lam), jit=False, name="RING")
+    ringnet = ml.base.ScaleX_FilterWrapper(ringnet)
+    ringnet = ml.base.LPF_FilterWrapper(
+        ringnet,
+        ml._LPF_CUTOFF_FREQ,
+        samp_freq=None if Ts is None else 1 / Ts,
+        quiet=True,
+    )
+    ringnet = ml.base.GroundTruthHeading_FilterWrapper(ringnet)
+    if add_Ts:
+        ringnet = ml.base.AddTs_FilterWrapper(ringnet, Ts)
+    return ringnet
+
+
+_TRAIN_TIMING_START = None
+_UNIQUE_ID = None
+
+
+def setup(
+    rr_joint_kwargs: None | dict = dict(),
+    rr_imp_joint_kwargs: None | dict = dict(),
+    suntay_joint_kwargs: None | dict = None,
+    train_timing_start: None | float = None,
+    unique_id: None | str = None,
+):
+    import time
+
+    from ring.algorithms import custom_joints
+
+    global _TRAIN_TIMING_START
+    global _UNIQUE_ID
+
+    if rr_joint_kwargs is not None:
+        custom_joints.register_rr_joint(**rr_joint_kwargs)
+
+    if rr_imp_joint_kwargs is not None:
+        custom_joints.register_rr_imp_joint(**rr_imp_joint_kwargs)
+
+    if suntay_joint_kwargs is not None:
+        custom_joints.register_suntay(**suntay_joint_kwargs)
+
+    if _TRAIN_TIMING_START is None:
+        _TRAIN_TIMING_START = time.time()
+
+    if train_timing_start is not None:
+        _TRAIN_TIMING_START = train_timing_start
+
+    if _UNIQUE_ID is None:
+        _UNIQUE_ID = hex(hash(time.time()))
+
+    if unique_id is not None:
+        _UNIQUE_ID = unique_id
+
+
+setup()

{imt_ring-1.3.13 → imt_ring-1.4.1}/src/ring/ml/__init__.py

@@ -13,28 +13,7 @@ from .optimizer import make_optimizer
 from .ringnet import RING
 from .train import train_fn
 
-_lpf_cutoff_freq = 10.0
-
-
-def RING_ICML24(params=None, eval: bool = True, **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)
-    if eval:
-        ringnet = base.LPF_FilterWrapper(ringnet, _lpf_cutoff_freq, samp_freq=None)
-    ringnet = base.GroundTruthHeading_FilterWrapper(ringnet)
-    return ringnet
+_LPF_CUTOFF_FREQ = 10.0
 
 
 def RNNO(
@@ -70,7 +49,7 @@ def RNNO(
     ringnet = base.NoGraph_FilterWrapper(ringnet, quat_normalize=return_quats)
     ringnet = base.ScaleX_FilterWrapper(ringnet)
     if eval and return_quats:
-        ringnet = base.LPF_FilterWrapper(ringnet, _lpf_cutoff_freq, samp_freq=samp_freq)
+        ringnet = base.LPF_FilterWrapper(ringnet, _LPF_CUTOFF_FREQ, samp_freq=samp_freq)
     if return_quats:
         ringnet = base.GroundTruthHeading_FilterWrapper(ringnet)
     return ringnet

{imt_ring-1.3.13 → imt_ring-1.4.1}/src/ring/ml/base.py

@@ -144,11 +144,13 @@ class LPF_FilterWrapper(AbstractFilterWrapper):
         cutoff_freq: float,
         samp_freq: float | None,
         filtfilt: bool = True,
+        quiet: bool = False,
         name="LPF_FilterWrapper",
     ) -> None:
         super().__init__(filter, name)
         self.samp_freq = samp_freq
         self._kwargs = dict(cutoff_freq=cutoff_freq, filtfilt=filtfilt)
+        self.quiet = quiet
 
     def apply(self, X, params=None, state=None, y=None, lam=None):
         if X.ndim == 4:
@@ -166,7 +168,7 @@ class LPF_FilterWrapper(AbstractFilterWrapper):
                 dt = X[0, 0, -1]
                 samp_freq = 1 / dt
 
-        if self.samp_freq is None:
+        if self.samp_freq is None and not self.quiet:
             print(f"Detected the following sampling rates from `X`: {samp_freq}")
 
         yhat, state = super().apply(X, params, state, y, lam)
@@ -290,3 +292,26 @@ class NoGraph_FilterWrapper(AbstractFilterWrapper):
             yhat = ring.maths.safe_normalize(yhat)
 
         return yhat, state
+
+
+class AddTs_FilterWrapper(AbstractFilterWrapper):
+    def __init__(
+        self, filter: AbstractFilter, Ts: float | None, name="AddTs_FilterWrapper"
+    ) -> None:
+        super().__init__(filter, name)
+        self.Ts = Ts
+
+    def _add_Ts(self, X):
+        if self.Ts is None:
+            assert X.shape[-1] == 10
+            return X
+        else:
+            assert X.shape[-1] == 9
+            X_Ts = jnp.ones(X.shape[:-1] + (1,)) * self.Ts
+            return jnp.concatenate((X, X_Ts), axis=-1)
+
+    def init(self, bs=None, X=None, lam=None, seed: int = 1):
+        return super().init(bs, self._add_Ts(X), lam, seed)
+
+    def apply(self, X, params=None, state=None, y=None, lam=None):
+        return super().apply(self._add_Ts(X), params, state, y, lam)

imt_ring-1.4.1/tests/test_quickstart_example.py

@@ -0,0 +1,22 @@
+import jax
+import numpy as np
+
+import ring
+
+
+def test_quickstart_exampe():
+    T: int = 30  # sequence length     [s]
+    Ts: float = 0.01  # sampling interval   [s]
+    B: int = 1  # batch size
+    lam: list[int] = [0, 1, 2]  # parent array
+    N: int = len(lam)  # number of bodies
+    T_i: int = int(T / Ts)  # number of timesteps
+
+    X = np.zeros((B, T_i, N, 9))
+
+    ringnet = ring.RING(lam, Ts)
+    yhat, state = ringnet.apply(X)
+    assert yhat.shape == (B, T_i, N, 4)
+    assert state["~"]["inner_cell_state"].shape == (B, N, 2, 400)
+
+    _ = jax.jit(ringnet.apply)(X, state=state)

imt_ring-1.3.13/src/ring/__init__.py

@@ -1,63 +0,0 @@
-from . import algebra
-from . import algorithms
-from . import base
-from . import io
-from . import maths
-from . import ml
-from . import rendering
-from . import sim2real
-from . import spatial
-from . import sys_composer
-from . import utils
-from .algorithms import join_motionconfigs
-from .algorithms import JointModel
-from .algorithms import MotionConfig
-from .algorithms import RCMG
-from .algorithms import register_new_joint_type
-from .algorithms import step
-from .base import State
-from .base import System
-from .base import Transform
-from .ml import RING
-
-_TRAIN_TIMING_START = None
-_UNIQUE_ID = None
-
-
-def setup(
-    rr_joint_kwargs: None | dict = dict(),
-    rr_imp_joint_kwargs: None | dict = dict(),
-    suntay_joint_kwargs: None | dict = None,
-    train_timing_start: None | float = None,
-    unique_id: None | str = None,
-):
-    import time
-
-    from ring.algorithms import custom_joints
-
-    global _TRAIN_TIMING_START
-    global _UNIQUE_ID
-
-    if rr_joint_kwargs is not None:
-        custom_joints.register_rr_joint(**rr_joint_kwargs)
-
-    if rr_imp_joint_kwargs is not None:
-        custom_joints.register_rr_imp_joint(**rr_imp_joint_kwargs)
-
-    if suntay_joint_kwargs is not None:
-        custom_joints.register_suntay(**suntay_joint_kwargs)
-
-    if _TRAIN_TIMING_START is None:
-        _TRAIN_TIMING_START = time.time()
-
-    if train_timing_start is not None:
-        _TRAIN_TIMING_START = train_timing_start
-
-    if _UNIQUE_ID is None:
-        _UNIQUE_ID = hex(hash(time.time()))
-
-    if unique_id is not None:
-        _UNIQUE_ID = unique_id
-
-
-setup()