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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.13
+Version: 1.4.0
 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.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ring/__init__.py,sha256=iNvbAZi7Qfa69IbL1z4lB7zHL8WusV5fBrKah2la-Gc,1566
+ring/__init__.py,sha256=-70A1E5LQQyGdOA_u9PlYZb1d5Fz0yXnvRoZOrbxq6o,3781
 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=16zdF72XLljeFRWaEr1O-9M5Sw2ppk5yvAoAPxA5EJU,9693
 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
@@ -79,7 +79,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.0.dist-info/METADATA,sha256=4o8Vsz-U7Ekp_4ipyhcCy1vlhCgVZwEZYbwLsfmzy34,3104
+imt_ring-1.4.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.4.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.4.0.dist-info/RECORD,,

ring/__init__.py

@@ -18,7 +18,74 @@ 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 > (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"
+        )
+
+    params = Path(__file__).parent.joinpath("ml/params/0x13e3518065c21cd8.pickle")
+
+    ringnet = ml.RING(params=params, lam=tuple(lam), jit=False)
+    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
+    )
+    ringnet = ml.base.GroundTruthHeading_FilterWrapper(ringnet)
+    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

@@ -290,3 +290,24 @@ 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=None) -> 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)