imt-ring 1.3.13__py3-none-any.whl → 1.4.1__py3-none-any.whl

{imt_ring-1.3.13.dist-info → imt_ring-1.4.1.dist-info}/METADATA

@@ -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.dist-info → imt_ring-1.4.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ring/__init__.py,sha256=iNvbAZi7Qfa69IbL1z4lB7zHL8WusV5fBrKah2la-Gc,1566
+ring/__init__.py,sha256=1kDuGEkC-Fkuzlijg4PFGg1r60HDmk1xsNNTiLKRToA,4331
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
 ring/base.py,sha256=YFPrUWelWswEhq8x8Byv-5pK64mipiGW6x5IlMr4we4,33803
 ring/maths.py,sha256=jJr_kr78-XDce8B4tXQ2Li-jBntVQhaS8csxglCsj8A,12193
@@ -50,8 +50,8 @@ ring/io/xml/from_xml.py,sha256=8b44sPVWgoY8JGJZLpJ8M_eLfcfu3IsMtBzSytPTPmw,9234
 ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,1562
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
 ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
-ring/ml/__init__.py,sha256=52LpEjni5lG-ov5-3ocodH-vKZxNcFMU7W9XfjDicp0,2113
-ring/ml/base.py,sha256=PQ72VasEqlecBZgWP5HE5rWYyLiLq7nCVLymXo9f0dw,8959
+ring/ml/__init__.py,sha256=8SZTCs9rJ1kzR0Psh7lUzFhIMhKRPIK41mVfxJAGyMo,1471
+ring/ml/base.py,sha256=-3JQ27zMFESNn5zeNer14GJU2yQgiqDcJUaULOeSyp8,9799
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
@@ -60,6 +60,7 @@ ring/ml/rnno_v1.py,sha256=T4SKG7iypqn2HBQLKhDmJ2Slj2Z5jtUBHvX_6aL8pyM,1103
 ring/ml/train.py,sha256=huUfMK6eotS6BRrQKoZ-AUG0um3jlqpfQFZNJT8LKiE,10854
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
+ring/ml/params/0x1d76628065a71e0f.pickle,sha256=YTNVuvfw-nCRD9BH1PZYcR9uCFpNWDhw8Lc50eDn_EE,9351038
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
 ring/rendering/base_render.py,sha256=s5dF-GVBqjiWkqVuPQMtTLuM7EtA-YrB7RVWFfIaQ1I,8956
 ring/rendering/mujoco_render.py,sha256=aluzQJp3jrDdPfAyNmQuXIHRfgfBTCCZQqxKOx_0D2s,7770
@@ -79,7 +80,7 @@ ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
 ring/utils/utils.py,sha256=mIcKNv5v2de8HrG7bAhl2bNfmwkMZyIIwFkJq2XWMOI,5357
-imt_ring-1.3.13.dist-info/METADATA,sha256=29lL2WaS8JaByX0qyjVm-OMC916Xs8U1A7pRyybHSrA,3105
-imt_ring-1.3.13.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.13.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.13.dist-info/RECORD,,
+imt_ring-1.4.1.dist-info/METADATA,sha256=uxnuzRi_YRN2-Toh7TyiglMt5O4bQ2H7TpWHq_-Yj7g,3104
+imt_ring-1.4.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.4.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.4.1.dist-info/RECORD,,

ring/__init__.py

@@ -18,7 +18,87 @@ from .algorithms import step
 from .base import State
 from .base import System
 from .base import Transform
-from .ml import RING
+
+
+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

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

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)