imt-ring 1.3.12__tar.gz → 1.4.0__tar.gz

{imt_ring-1.3.12 → imt_ring-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.12
+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.12 → imt_ring-1.4.0}/pyproject.toml

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

{imt_ring-1.3.12 → imt_ring-1.4.0}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.12
+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.12 → imt_ring-1.4.0}/src/imt_ring.egg-info/SOURCES.txt

@@ -99,6 +99,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.0/src/ring/__init__.py

@@ -0,0 +1,130 @@
+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 > (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
+
+
+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.12 → imt_ring-1.4.0}/src/ring/algorithms/generator/batch.py

@@ -154,25 +154,20 @@ def _data_fn_from_paths(
 
     # expanduser
     paths = [utils.parse_path(p, mkdir=False) for p in paths]
-
     extensions = list(set([Path(p).suffix for p in paths]))
     assert len(extensions) == 1, f"{extensions}"
+    h5 = extensions[0] == ".h5"
 
-    if extensions[0] == ".h5":
-        N = sum([utils.hdf5_load_length(p) for p in paths])
-
-    if extensions[0] == ".h5" and not load_all_into_memory:
+    if h5 and not load_all_into_memory:
 
         def data_fn(indices: list[int]):
             tree = utils.hdf5_load_from_multiple(paths, indices)
             return tree if tree_transform is None else tree_transform(tree)
 
+        N = sum([utils.hdf5_load_length(p) for p in paths])
     else:
 
-        if extensions[0] == ".h5":
-            load_from_path = utils.hdf5_load
-        else:
-            load_from_path = utils.pickle_load
+        load_from_path = utils.hdf5_load if h5 else utils.pickle_load
 
         def load_fn(path):
             tree = load_from_path(path)
@@ -190,8 +185,10 @@ def _data_fn_from_paths(
                 _list_of_data += load_fn(p)
 
         N = len(_list_of_data)
-
-        list_of_data = _replace_elements_w_nans(_list_of_data, include_samples)
+        list_of_data = _replace_elements_w_nans(
+            _list_of_data,
+            include_samples if include_samples is not None else list(range(N)),
+        )
 
         if include_samples is not None:
             list_of_data = [

{imt_ring-1.3.12 → imt_ring-1.4.0}/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.12 → imt_ring-1.4.0}/src/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)

imt_ring-1.4.0/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.12/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()