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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.3
+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.3.dist-info → imt_ring-1.6.4.dist-info}/RECORD

@@ -1,11 +1,11 @@
-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=06x7RfhoQ6dx1B_TAEuCKxNTiicQDDBxcmzRtsCAxsM,18125
@@ -14,7 +14,7 @@ ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwa
 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=ZW6tJSUiNKtRCYIsbCf6e6kRBbmuCLzbZj4ppNLVwJY,14368
+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=LUw1Wc2YrmMRRh4RF704ob3bZOXktAZAbbLoBm_p1yw,9131
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
@@ -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.3.dist-info/METADATA,sha256=uFTNWR0YQbQLE50Oby-D2H3NSpdycgCwzKaz_UxxSP8,3104
-imt_ring-1.6.3.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-imt_ring-1.6.3.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.3.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
@@ -47,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
@@ -141,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))
@@ -170,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(

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