imt-ring 1.6.11__tar.gz → 1.6.12__tar.gz

{imt_ring-1.6.11 → imt_ring-1.6.12}/PKG-INFO

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

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

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/imt_ring.egg-info/PKG-INFO

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

@@ -84,6 +84,7 @@ src/ring/utils/__init__.py
 src/ring/utils/backend.py
 src/ring/utils/batchsize.py
 src/ring/utils/colab.py
+src/ring/utils/dataloader.py
 src/ring/utils/hdf5.py
 src/ring/utils/normalizer.py
 src/ring/utils/path.py

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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:

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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]:

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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.

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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(

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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
         )

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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
@@ -190,36 +189,30 @@ def unique_id() -> str:
 
 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)

imt_ring-1.6.12/src/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

{imt_ring-1.6.11 → imt_ring-1.6.12}/src/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)