imt-ring 1.6.11__py3-none-any.whl → 1.6.13__py3-none-any.whl

{imt_ring-1.6.11.dist-info → imt_ring-1.6.13.dist-info}/METADATA

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

@@ -1,22 +1,22 @@
 ring/__init__.py,sha256=k7tL-XgggUwWxHCXyv60rQn-OcXHPg82QcIUkKLEd-c,5057
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=MkkziQx01sdMOpB8MFUDFgFlZUrXCFjpb8hS9yKHUyM,33751
+ring/base.py,sha256=Ystn1EjTyOXBhVm5koroV_YPUYtFxrteJLd-XR3kEL8,33840
 ring/maths.py,sha256=qPHH6TpHCK3TgExI98gNEySoSRKOwteN9McUlyUFipI,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=GOedL1STj6oXcXgMA7dB4PabvCQxPBbirJQhXBRuKqE,10929
-ring/algorithms/jcalc.py,sha256=bM8VARgqEiVPy7632geKYGk4MZddZfI8XHdW5kXF3HI,28594
+ring/algorithms/jcalc.py,sha256=bwfVH3qKEnUs6RFgEEeUBnecpBt-nf8cesJbNGDrE7g,28974
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
-ring/algorithms/sensors.py,sha256=0zLjd9TzOkU5W5GJU6Dk4QwYjwqs9AUlzUKU8aSX_dc,18126
+ring/algorithms/sensors.py,sha256=0xOzdQIc1kBF0CkoPXWWCx3MmV4SG3wj7knVnnMWq9M,18124
 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=LRAKxzrwq6fp4lgVw6IUg4i7isx3iqJLHvpFK1aTRcg,15732
+ring/algorithms/generator/base.py,sha256=vxUdA0ZeSNH3SOanL51qVRvCiJrmsWQyQX0g2fdm3Rg,15825
 ring/algorithms/generator/batch.py,sha256=9yFxVv11hij-fJXGPxA3zEh1bE2_jrZk0R7kyGaiM5c,2551
-ring/algorithms/generator/finalize_fns.py,sha256=LUw1Wc2YrmMRRh4RF704ob3bZOXktAZAbbLoBm_p1yw,9131
+ring/algorithms/generator/finalize_fns.py,sha256=559sGXs06n46p-eme0SE8hn0lXwGT0P2r3-52ElTldo,9861
 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
@@ -53,10 +53,10 @@ ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
 ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
-ring/ml/ml_utils.py,sha256=GooyH5uxA6cJM7ZcWDUfSkSKq6dg7kCIbhkbjJs_rLw,6674
+ring/ml/ml_utils.py,sha256=1GXJfeoXbwCbRdYA2np3CbJpSupaw4eyf3quh9y4BO0,6462
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=Tb2WJ_cc5L3mk1lo0NOfkpXIzJZXf4PJ5aLPtHQyUmY,8650
-ring/ml/rnno_v1.py,sha256=T4SKG7iypqn2HBQLKhDmJ2Slj2Z5jtUBHvX_6aL8pyM,1103
+ring/ml/rnno_v1.py,sha256=ujyIkDxMSTag9iRFEmoHqfqSrlOFjcZs9_rBbLd8p9Q,1380
 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
@@ -76,14 +76,15 @@ ring/utils/__init__.py,sha256=MHHavc8YfjBlmB-zAV42QEQS_ebW7cy0lhWXEVyQU7s,720
 ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
 ring/utils/batchsize.py,sha256=FbOii7MDP4oPZd9GJOKehFatfnb6WZ0b9z349iZYs1A,1786
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
+ring/utils/dataloader.py,sha256=2CcsbUY2AZs8LraS5HTJXlEseuF-1gKmfyBkSsib-tE,3748
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=zRPfxYNesvgefkddd26oar6f9433LkMGkhp9dF3rPUs,1926
 ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
-ring/utils/utils.py,sha256=oGC7kh19s5zvmnUvWy8B3fBl9loVU58ppz91osk2m3w,6550
+ring/utils/utils.py,sha256=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
 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.11.dist-info/METADATA,sha256=kkQfOD5LOSzB4lR7LvkHeck6fB_KPNrSKIsvPizJAKI,3821
-imt_ring-1.6.11.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-imt_ring-1.6.11.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.11.dist-info/RECORD,,
+imt_ring-1.6.13.dist-info/METADATA,sha256=wMwfHX8PsYaxXZRldgqd71fGltIfUo_W9xVE9DSz5o0,3821
+imt_ring-1.6.13.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+imt_ring-1.6.13.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.13.dist-info/RECORD,,

{imt_ring-1.6.11.dist-info → imt_ring-1.6.13.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ring/algorithms/generator/base.py

@@ -321,6 +321,9 @@ def _build_mconfig_batched_generator(
                 "using the `randomize_motion_artifacts` flag, so it must be enabled."
             )
 
+    if dynamic_simulation:
+        finalize_fns.DynamicalSimulation.assert_test_system(sys)
+
     def _setup_fn(key: types.PRNGKey, sys: base.System) -> base.System:
         pipe = []
         if imu_motion_artifacts and randomize_motion_artifacts:

ring/algorithms/generator/finalize_fns.py

@@ -180,6 +180,23 @@ class DynamicalSimulation:
         self.overwrite_q_ref = overwrite_q_ref
         self.unroll_kwargs = unroll_kwargs
 
+    @staticmethod
+    def assert_test_system(sys: base.System) -> None:
+        "test that system has no zero mass bodies and no joints without damping"
+
+        def f(_, __, n, m, d):
+            assert d.size == 0 or m > 0, (
+                "Dynamic simulation is set to `True` which requires masses >= 0, "
+                f"but found body `{n}` with mass={float(m[0])}. This can lead to NaNs."
+            )
+
+            assert d.size == 0 or all(d > 0.0), (
+                "Dynamic simulation is set to `True` which requires dampings > 0, "
+                f"but found body `{n}` with damping={d}. This can lead to NaNs."
+            )
+
+        sys.scan(f, "lld", sys.link_names, sys.links.inertia.mass, sys.link_damping)
+
     def __call__(
         self, Xy: types.Xy, extras: types.OutputExtras
     ) -> tuple[types.Xy, types.OutputExtras]:

ring/algorithms/jcalc.py

@@ -205,7 +205,7 @@ def _is_feasible_config1(c: MotionConfig) -> bool:
             return False
         return True
 
-    return all(
+    cond1 = all(
         [
             dx_deltax_check(*args)
             for args in zip(
@@ -217,6 +217,15 @@ def _is_feasible_config1(c: MotionConfig) -> bool:
         ]
     )
 
+    # this one tests that the initial value is inside the feasible value range
+    # so e.g. if you choose pos0_min=-10 then you can't choose pos_min=-1
+    def inside_box_checks(x_min, x_max, x0_min, x0_max) -> bool:
+        return (x0_min >= x_min) and (x0_max <= x_max)
+
+    cond2 = inside_box_checks(c.pos_min, c.pos_max, c.pos0_min, c.pos0_max)
+
+    return cond1 and cond2
+
 
 def _find_interval(t: jax.Array, boundaries: jax.Array):
     """Find the interval of `boundaries` between which `t` lies.

ring/algorithms/sensors.py

@@ -131,7 +131,7 @@ def magnetometer(rot: jax.Array, magvec: jax.Array) -> jax.Array:
 # - gyr: rad/s
 # - mag: a.u.
 NOISE_LEVELS = {"acc": 0.048, "gyr": jnp.deg2rad(0.7), "mag": 0.01}
-BIAS_LEVELS = {"acc": 0.5, "gyr": jnp.deg2rad(3.6), "mag": 0.0}
+BIAS_LEVELS = {"acc": 0.5, "gyr": jnp.deg2rad(3), "mag": 0.0}
 
 
 def add_noise_bias(

ring/base.py

@@ -690,7 +690,9 @@ class System(_Base):
         transparent_segment_to_root: bool = True,
         **kwargs,
     ):
-        "`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
+        """`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent.
+        Note that the body in yhat that connects to -1, is parent-to-child!
+        """
         return ring.rendering.render_prediction(
             self, xs, yhat, transparent_segment_to_root, **kwargs
         )

ring/ml/ml_utils.py

@@ -12,7 +12,6 @@ import numpy as np
 from tree_utils import PyTree
 
 import ring
-from ring.utils import import_lib
 import wandb
 
 # An arbitrarily nested dictionary with Array leaves; Or strings
@@ -185,41 +184,37 @@ def on_cluster() -> bool:
 
 
 def unique_id() -> str:
+    if wandb.run is not None:
+        wandb.config.setdefault("unique_id", ring._UNIQUE_ID)
     return ring._UNIQUE_ID
 
 
 def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     from jax.experimental import jax2tf
+    import tensorflow as tf
 
-    tf = import_lib("tensorflow", "the function `save_model_tf`")
-
-    def _create_module(jax_func, input):
-        signature = jax.tree_map(
-            lambda arr: tf.TensorSpec(list(arr.shape), tf.float32), input
-        )
+    signature = jax.tree_map(
+        lambda arr: tf.TensorSpec(list(arr.shape), tf.float32), input
+    )
 
-        class RingTFModule(tf.Module):
-            def __init__(self, jax_func):
-                super().__init__()
-                self.tf_func = jax2tf.convert(jax_func, with_gradient=False)
+    tf_func = jax2tf.convert(jax_func, with_gradient=False)
 
-            @partial(
-                tf.function,
-                autograph=False,
-                jit_compile=True,
-                input_signature=signature,
-            )
-            def __call__(self, *args):
-                return self.tf_func(*args)
+    class RingTFModule(tf.Module):
+        @partial(
+            tf.function, autograph=False, jit_compile=True, input_signature=signature
+        )
+        def __call__(self, *args):
+            return tf_func(*args)
 
-        return RingTFModule(jax_func)
+    model = RingTFModule()
 
-    model = _create_module(jax_func, input)
     tf.saved_model.save(
         model,
         path,
         options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
+        signatures={"default": model.__call__},
     )
+
     if validate:
         output_jax = jax_func(*input)
         output_tf = tf.saved_model.load(path)(*input)

ring/ml/rnno_v1.py

@@ -4,6 +4,8 @@ import haiku as hk
 import jax
 import jax.numpy as jnp
 
+from .ringnet import LSTM
+
 
 def rnno_v1_forward_factory(
     output_dim: int,
@@ -13,18 +15,29 @@ def rnno_v1_forward_factory(
     act_fn_linear=jax.nn.relu,
     act_fn_rnn=jax.nn.elu,
     lam: Optional[tuple[int]] = None,
+    celltype: str = "gru",
 ):
     # unused
     del lam
 
+    if celltype == "gru":
+        _cell = hk.GRU
+        _factor = 1
+    elif celltype == "lstm":
+        _cell = LSTM
+        _factor = 2
+    else:
+        raise NotImplementedError
+
     @hk.without_apply_rng
     @hk.transform_with_state
     def forward_fn(X):
         assert X.shape[-2] == 1
 
         for i, n_units in enumerate(rnn_layers):
+            n_units = _factor * n_units
             state = hk.get_state(f"rnn_{i}", shape=[1, n_units], init=jnp.zeros)
-            X, state = hk.dynamic_unroll(hk.GRU(n_units), X, state)
+            X, state = hk.dynamic_unroll(_cell(n_units), X, state)
             hk.set_state(f"rnn_{i}", state)
 
             if layernorm:

ring/utils/dataloader.py

@@ -0,0 +1,159 @@
+import os
+import random
+from typing import Callable, Optional
+
+import jax
+import numpy as np
+from ring.utils import parse_path
+from ring.utils import pickle_load
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+import tqdm
+from tree_utils import PyTree
+
+
+def make_generator(
+    *paths,
+    batch_size,
+    transform,
+    shuffle=True,
+    seed: int = 1,
+    backend: str = "eager",
+    **kwargs,
+):
+    if backend == "grain":
+        _make_gen = pygrain_generator
+    elif backend == "torch":
+        _make_gen = pytorch_generator
+    elif backend == "eager":
+        _make_gen = eager_generator
+    else:
+        raise NotImplementedError
+
+    return _make_gen(
+        *paths,
+        batch_size=batch_size,
+        transform=transform,
+        shuffle=shuffle,
+        seed=seed,
+        **kwargs,
+    )
+
+
+T = PyTree[np.ndarray]
+
+
+class _Dataset(Dataset):
+    def __init__(self, *paths, transform):
+
+        self.files = [self.listdir(path) for path in paths]
+        Ns = set([len(f) for f in self.files])
+        assert len(Ns) == 1, f"{Ns}"
+
+        self.P = len(self.files)
+        self.N = list(Ns)[0]
+        self.transform = transform
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, idx: int):
+        element = [pickle_load(self.files[p][idx]) for p in range(self.P)]
+        if self.transform is not None:
+            element = self.transform(element)
+        return element
+
+    @staticmethod
+    def listdir(path: str) -> list:
+        return [parse_path(path, file) for file in os.listdir(path)]
+
+    def __call__(self, idx: int):
+        return self[idx]
+
+
+class TransformTransform:
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, element):
+        if self.transform is None:
+            return element
+        return self.transform(element, np.random.default_rng())
+
+
+def pytorch_generator(
+    *paths,
+    batch_size: int,
+    transform: Optional[Callable[[T], T]] = None,
+    shuffle=True,
+    seed: int = 1,
+    **kwargs,
+):
+    torch.manual_seed(seed)
+
+    ds = _Dataset(*paths, transform=TransformTransform(transform))
+    dl = DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        multiprocessing_context="spawn" if kwargs.get("num_workers", 0) > 0 else None,
+        **kwargs,
+    )
+    dl_iter = iter(dl)
+
+    def to_numpy(tree: PyTree[torch.Tensor]):
+        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+
+    def generator(_):
+        nonlocal dl, dl_iter
+        try:
+            return to_numpy(next(dl_iter))
+        except StopIteration:
+            dl_iter = iter(dl)
+            return to_numpy(next(dl_iter))
+
+    return generator
+
+
+def eager_generator(
+    *paths,
+    batch_size: int,
+    transform: Optional[Callable[[T], T]] = None,
+    shuffle=True,
+    seed=1,
+):
+    from ring import RCMG
+
+    random.seed(seed)
+
+    ds = _Dataset(*paths, transform=TransformTransform(transform))
+    data = [ds[i] for i in tqdm.tqdm(range(len(ds)), total=len(ds))]
+    return RCMG.eager_gen_from_list(data, batch_size, shuffle=shuffle)
+
+
+def pygrain_generator(
+    *paths, batch_size: int, transform=None, shuffle=True, seed=1, **kwargs
+):
+
+    import grain.python as pygrain  # type: ignore
+
+    class _Transform(pygrain.RandomMapTransform):
+        def random_map(self, element, rng: np.random.Generator):
+            return transform(element, rng)
+
+    ds = _Dataset(*paths, transform=None)
+    dl = pygrain.load(
+        ds,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        seed=seed,
+        transformations=[_Transform()],
+        **kwargs,
+    )
+    iter_dl = iter(dl)
+
+    def generator(_):
+        return next(iter_dl)
+
+    return generator

ring/utils/utils.py

@@ -3,6 +3,7 @@ import io
 import pickle
 import random
 from typing import Optional
+import warnings
 
 import jax
 import jax.numpy as jnp
@@ -195,7 +196,7 @@ def replace_elements_w_nans(
     assert min(include_elements) >= 0
     assert max(include_elements) < len(list_of_data)
 
-    def _is_nan(ele: tree_utils.PyTree, i: int):
+    def _is_nan(ele: tree_utils.PyTree, i: int, verbose: bool):
         isnan = np.any(
             [np.any(np.isnan(arr)) for arr in jax.tree_util.tree_leaves(ele)]
         )
@@ -205,13 +206,22 @@ def replace_elements_w_nans(
             return True
         return False
 
+    list_of_isnan = [int(_is_nan(e, 0, False)) for e in list_of_data]
+    perc_of_isnan = sum(list_of_isnan) / len(list_of_data)
+
+    if perc_of_isnan >= 0.02:
+        warnings.warn(
+            f"{perc_of_isnan * 100}% of {len(list_of_data)} datapoints are NaN"
+        )
+        assert perc_of_isnan != 1
+
     list_of_data_nonan = []
     for i, ele in enumerate(list_of_data):
-        if _is_nan(ele, i):
+        if _is_nan(ele, i, verbose):
             while True:
                 j = random.choice(include_elements)
                 ele_j = list_of_data[j]
-                if not _is_nan(ele_j, j):
+                if not _is_nan(ele_j, j, verbose):
                     ele = pytree_deepcopy(ele_j)
                     break
         list_of_data_nonan.append(ele)