imt-ring 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl

{imt_ring-1.6.2.dist-info → imt_ring-1.6.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.2
+Version: 1.6.4
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring
@@ -60,6 +60,36 @@ pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-re
 
 Available [here](https://simipixel.github.io/ring/).
 
+## Quickstart Example
+```python
+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)
+```
+
 ### Known fixes
 
 #### Offscreen rendering with Mujoco

{imt_ring-1.6.2.dist-info → imt_ring-1.6.4.dist-info}/RECORD

@@ -1,22 +1,22 @@
-ring/__init__.py,sha256=2v6WHlNPucj1XGhDYw-3AlMQGTqH-e4KYK0IaMnBV5s,4760
+ring/__init__.py,sha256=k7tL-XgggUwWxHCXyv60rQn-OcXHPg82QcIUkKLEd-c,5057
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=kzBQ54V2xq4KsqRzflyMQ64V-jl8j7eIAsIPIE0gFDk,33127
+ring/base.py,sha256=BGAJE3PSOUnTHte4UesJc1J7MQraIEiVpStkhrgXhaI,33245
 ring/maths.py,sha256=zGm5XagiKTaIJp310VcqVEVUuLhv3FPS-TJ-TFzIrwM,12207
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
 ring/algorithms/_random.py,sha256=fc26yEQjSjtf0NluZ41CyeGIRci0ldrRlThueHR9H7U,14007
-ring/algorithms/dynamics.py,sha256=_TwclBXe6vi5C5iJWAIeUIJEIMHQ_1QTmnHvCEpVO0M,10867
+ring/algorithms/dynamics.py,sha256=GOedL1STj6oXcXgMA7dB4PabvCQxPBbirJQhXBRuKqE,10929
 ring/algorithms/jcalc.py,sha256=bM8VARgqEiVPy7632geKYGk4MZddZfI8XHdW5kXF3HI,28594
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
-ring/algorithms/sensors.py,sha256=QSIcU_sEB_tRo-ADD_66ZD01LMmJlSG0op6YnM-Gai8,17965
+ring/algorithms/sensors.py,sha256=06x7RfhoQ6dx1B_TAEuCKxNTiicQDDBxcmzRtsCAxsM,18125
 ring/algorithms/custom_joints/__init__.py,sha256=fzeE7TdUhmGgbbFAyis1tKcyQ4Fo8LigDwD3hUVnH_w,316
 ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwaMQyUjYv7EtsPiU3A,2402
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
 ring/algorithms/custom_joints/suntay.py,sha256=tOEGM304XciHO4pmvxr4faA4xXVO4N2HlPdFmXKbcrw,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=KQSg9uhhR-rC563busVFx4gJrqOx3BXdaChozO9gwTA,14224
+ring/algorithms/generator/base.py,sha256=00789VeBrYcKx5BjsiGu-d4UzO6FGRz-YUKiiLUOL2Q,14497
 ring/algorithms/generator/batch.py,sha256=ylootnXmj-JyuB_f5OCknHst9wFKO3gkjQbMrFNXY2g,2513
-ring/algorithms/generator/finalize_fns.py,sha256=L_5wIVA7g0P4P2U6EmgcvsoI-YuF3TOaHBwk5_oEaUU,9077
+ring/algorithms/generator/finalize_fns.py,sha256=LUw1Wc2YrmMRRh4RF704ob3bZOXktAZAbbLoBm_p1yw,9131
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
 ring/algorithms/generator/setup_fns.py,sha256=MFz3czHBeWs1Zk1A8O02CyQpQ-NCyW9PMpbqmKit6es,1455
@@ -55,7 +55,7 @@ ring/ml/base.py,sha256=-3JQ27zMFESNn5zeNer14GJU2yQgiqDcJUaULOeSyp8,9799
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 ring/ml/ml_utils.py,sha256=GooyH5uxA6cJM7ZcWDUfSkSKq6dg7kCIbhkbjJs_rLw,6674
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
-ring/ml/ringnet.py,sha256=rgje5AKUKpT8K-vbE9_SgZ3IijR8TJEHnaqxsE57Mhc,8617
+ring/ml/ringnet.py,sha256=Tb2WJ_cc5L3mk1lo0NOfkpXIzJZXf4PJ5aLPtHQyUmY,8650
 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
@@ -83,7 +83,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=k7t-QxMWrNRnjfNB9rSobmLCmhJigE8__gkT-Il0Ee4,6492
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.2.dist-info/METADATA,sha256=S_LlVrmdRPQCzT5aeRoSWyOQ3eBJBL1D33tUgXMUEso,3104
-imt_ring-1.6.2.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-imt_ring-1.6.2.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.2.dist-info/RECORD,,
+imt_ring-1.6.4.dist-info/METADATA,sha256=3uiX-NEHrXZ4QurQPYQUPRlul7aVLBZYIr0w17Yo54E,3922
+imt_ring-1.6.4.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+imt_ring-1.6.4.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.4.dist-info/RECORD,,

ring/__init__.py

@@ -20,52 +20,55 @@ from .base import System
 from .base import Transform
 
 
-def RING(lam: list[int] | None, Ts: float | None, **kwargs):
+def RING(lam: list[int] | None, Ts: float | None, **kwargs) -> ml.AbstractFilter:
     """Creates the RING network.
 
     Params:
         lam: parent array, if `None` must be given via `ringnet.apply(..., lam=lam)`
         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)
-    >>> # yhat[b, :, i] is the orientation from body `i` to parent body `lam[i]`
-    >>>
-    >>> # 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)
-
+    Returns:
+        ring.ml.AbstractFilter: An instantiation of `ring.ml.ringnet.RING` with trained
+            parameters.
+
+    Examples:
+        >>> 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)
+        >>> # yhat[b, :, i] is the orientation from body `i` to parent body `lam[i]`
+        >>>
+        >>> # 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

ring/algorithms/dynamics.py

@@ -303,6 +303,7 @@ def step(
     taus: Optional[jax.Array] = None,
     n_substeps: int = 1,
 ) -> base.State:
+    "Steps the dynamics. Returns the state of next timestep."
     assert sys.q_size() == state.q.size
     if taus is None:
         taus = jnp.zeros_like(state.qd)

ring/algorithms/generator/base.py

@@ -4,6 +4,7 @@ import warnings
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 import tree_utils
 
 from ring import base
@@ -30,6 +31,7 @@ class RCMG:
         add_X_jointaxes_kwargs: dict = dict(),
         add_y_relpose: bool = False,
         add_y_rootincl: bool = False,
+        add_y_rootincl_kwargs: dict = dict(),
         sys_ml: Optional[base.System] = None,
         randomize_positions: bool = False,
         randomize_motion_artifacts: bool = False,
@@ -46,6 +48,7 @@ class RCMG:
         cor: bool = False,
         disable_tqdm: bool = False,
     ) -> None:
+        "Random Chain Motion Generator"
 
         sys, config = utils.to_list(sys), utils.to_list(config)
         sys_ml = sys[0] if sys_ml is None else sys_ml
@@ -67,6 +70,7 @@ class RCMG:
                     add_X_jointaxes_kwargs=add_X_jointaxes_kwargs,
                     add_y_relpose=add_y_relpose,
                     add_y_rootincl=add_y_rootincl,
+                    add_y_rootincl_kwargs=add_y_rootincl_kwargs,
                     sys_ml=sys_ml,
                     randomize_positions=randomize_positions,
                     randomize_motion_artifacts=randomize_motion_artifacts,
@@ -139,7 +143,9 @@ class RCMG:
 
         return n_calls
 
-    def to_list(self, sizes: int | list[int] = 1, seed: int = 1):
+    def to_list(
+        self, sizes: int | list[int] = 1, seed: int = 1
+    ) -> list[tree_utils.PyTree[np.ndarray]]:
         "Returns list of unbatched sequences as numpy arrays."
         repeats = self._compute_repeats(sizes)
         sizes = list(jnp.array(repeats) * jnp.array(self._size_of_generators))
@@ -168,7 +174,7 @@ class RCMG:
         seed: int = 1,
         overwrite: bool = True,
     ) -> None:
-        data = tree_utils.tree_batch(self.to_list(sizes, seed))
+        data = tree_utils.tree_batch(self.to_list(sizes, seed), backend="numpy")
         utils.pickle_save(data, path, overwrite=overwrite)
 
     def to_eager_gen(
@@ -232,6 +238,7 @@ def _build_mconfig_batched_generator(
     add_X_jointaxes_kwargs: dict,
     add_y_relpose: bool,
     add_y_rootincl: bool,
+    add_y_rootincl_kwargs: dict,
     sys_ml: base.System,
     randomize_positions: bool,
     randomize_motion_artifacts: bool,

ring/algorithms/generator/finalize_fns.py

@@ -77,12 +77,13 @@ class RelPose:
 
 
 class RootIncl:
-    def __init__(self, sys: base.System):
+    def __init__(self, sys: base.System, **kwargs):
         self.sys = sys
+        self.kwargs = kwargs
 
     def __call__(self, Xy, extras):
         (X, y), (key, q, x, sys_x) = Xy, extras
-        y_root_incl = sensors.root_incl(self.sys, x, sys_x)
+        y_root_incl = sensors.root_incl(self.sys, x, sys_x, **self.kwargs)
         y = utils.dict_union(y, y_root_incl)
         return (X, y), (key, q, x, sys_x)
 

ring/algorithms/sensors.py

@@ -330,7 +330,10 @@ def rel_pose(
 
 
 def root_incl(
-    sys: base.System, x: base.Transform, sys_x: base.System
+    sys: base.System,
+    x: base.Transform,
+    sys_x: base.System,
+    child_to_parent: bool = False,
 ) -> dict[str, jax.Array]:
     # (time, nlinks, 4) -> (nlinks, time, 4)
     rots = x.rot.transpose((1, 0, 2))
@@ -341,8 +344,10 @@ def root_incl(
     def f(_, __, name: str, parent: int):
         if parent != -1:
             return
-        q_eps_to_i = maths.quat_project(rots[l_map[name]], jnp.array([0.0, 0, 1]))[1]
-        y[name] = maths.quat_inv(q_eps_to_i)
+        q_i = maths.quat_project(rots[l_map[name]], jnp.array([0.0, 0, 1]))[1]
+        if child_to_parent:
+            q_i = maths.quat_inv(q_i)
+        y[name] = q_i
 
     sys.scan(f, "ll", sys.link_names, sys.link_parents)
 
@@ -350,7 +355,10 @@ def root_incl(
 
 
 def root_full(
-    sys: base.System, x: base.Transform, sys_x: base.System
+    sys: base.System,
+    x: base.Transform,
+    sys_x: base.System,
+    child_to_parent: bool = False,
 ) -> dict[str, jax.Array]:
     # (time, nlinks, 4) -> (nlinks, time, 4)
     rots = x.rot.transpose((1, 0, 2))
@@ -361,8 +369,10 @@ def root_full(
     def f(_, __, name: str, parent: int):
         if parent != -1:
             return
-        q_eps_to_i = rots[l_map[name]]
-        y[name] = maths.quat_inv(q_eps_to_i)
+        q_i = rots[l_map[name]]
+        if child_to_parent:
+            q_i = maths.quat_inv(q_i)
+        y[name] = q_i
 
     sys.scan(f, "ll", sys.link_names, sys.link_parents)
 

ring/base.py

@@ -113,7 +113,9 @@ class _Base:
 class Transform(_Base):
     """Represents the Transformation from Plücker A to Plücker B,
     where B is located relative to A at `pos` in frame A and `rot` is the
-    relative quaternion from A to B."""
+    relative quaternion from A to B.
+    Create using `Transform.create(pos=..., rot=...)
+    """
 
     pos: Vector
     rot: Quaternion
@@ -399,6 +401,7 @@ QD_WIDTHS = {
 
 @struct.dataclass
 class System(_Base):
+    "System object. Create using `System.create(path_xml)`"
     link_parents: list[int] = struct.field(False)
     links: Link
     link_types: list[str] = struct.field(False)

ring/ml/ringnet.py

@@ -200,6 +200,7 @@ class RING(ml_base.AbstractFilter):
         forward_factory=make_ring,
         **kwargs,
     ):
+        "Untrained RING network"
         self.forward_lam_factory = partial(forward_factory, **kwargs)
         self.params = self._load_params(params)
         self.lam = lam