imt-ring 1.6.31__tar.gz → 1.6.33__tar.gz

{imt_ring-1.6.31 → imt_ring-1.6.33}/PKG-INFO

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

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

{imt_ring-1.6.31 → imt_ring-1.6.33}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.31
+Version: 1.6.33
 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.31 → imt_ring-1.6.33}/src/ring/algorithms/generator/base.py

@@ -213,6 +213,8 @@ class RCMG:
                 )
                 save_fn(d, file)
                 i += 1
+            # cleanup
+            del data
 
         gens, n_calls = self._generators_ncalls(sizes)
         batch.generators_eager(gens, n_calls, callback, seed, self._disable_tqdm)

{imt_ring-1.6.31 → imt_ring-1.6.33}/src/ring/algorithms/generator/batch.py

@@ -1,3 +1,4 @@
+import gc
 from typing import Callable
 
 import jax
@@ -83,4 +84,8 @@ def generators_eager(
 
             sample_flat, _ = jax.tree_util.tree_flatten(sample)
             size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
-            callback([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+            callback([jax.tree_map(lambda a: a[i].copy(), sample) for i in range(size)])
+
+            # cleanup
+            del sample, sample_flat
+            gc.collect()

{imt_ring-1.6.31 → imt_ring-1.6.33}/src/ring/algorithms/jcalc.py

@@ -412,6 +412,9 @@ def _find_interval(t: jax.Array, boundaries: jax.Array):
 def join_motionconfigs(
     configs: list[MotionConfig], boundaries: list[float]
 ) -> MotionConfig:
+    # to avoid a circular import due to `ring.utils.randomize_sys` importing `jcalc`
+    from ring.utils import tree_equal
+
     assert len(configs) == (
         len(boundaries) + 1
     ), "length of `boundaries` should be one less than length of `configs`"
@@ -434,10 +437,22 @@ def join_motionconfigs(
     time_independent_fields = [key for key in attrs if not is_time_dependent_field(key)]
 
     for time_dep_field in time_independent_fields:
-        field_values = set([getattr(config, time_dep_field) for config in configs])
-        assert (
-            len(field_values) == 1
-        ), f"MotionConfig.{time_dep_field}={field_values}. Should be one unique value.."
+        try:
+            field_values = set([getattr(config, time_dep_field) for config in configs])
+            assert (
+                len(field_values) == 1
+            ), f"MotionConfig.{time_dep_field}={field_values}. "
+            "Should be one unique value.."
+        except (
+            TypeError
+        ):  # dict is not hashable so test equality of all elements differently
+            comparison_ele = getattr(configs[0], time_dep_field)
+            for other_config in configs[1:]:
+                other_ele = getattr(other_config, time_dep_field)
+                assert tree_equal(
+                    comparison_ele, other_ele
+                ), f"MotionConfig.{time_dep_field} with {comparison_ele} != {other_ele}"
+                " Should be one unique value.."
 
     changes = {field: new_value(field) for field in time_dependent_fields}
     return replace(configs[0], **changes)

{imt_ring-1.6.31 → imt_ring-1.6.33}/src/ring/utils/dataloader_torch.py

@@ -1,8 +1,9 @@
 import os
-from typing import Optional
+from typing import Any, Optional
 import warnings
 
 import jax
+import numpy as np
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
@@ -12,7 +13,7 @@ from ring.utils import parse_path
 from ring.utils import pickle_load
 
 
-class FolderOfPickleFilesDataset(Dataset):
+class FolderOfFilesDataset(Dataset):
     def __init__(self, path, transform=None):
         self.files = self.listdir(path)
         self.transform = transform
@@ -22,7 +23,7 @@ class FolderOfPickleFilesDataset(Dataset):
         return self.N
 
     def __getitem__(self, idx: int):
-        element = pickle_load(self.files[idx])
+        element = self._load_file(self.files[idx])
         if self.transform is not None:
             element = self.transform(element)
         return element
@@ -31,6 +32,10 @@ class FolderOfPickleFilesDataset(Dataset):
     def listdir(path: str) -> list:
         return [parse_path(path, file) for file in os.listdir(path)]
 
+    @staticmethod
+    def _load_file(file_path: str) -> Any:
+        return pickle_load(file_path)
+
 
 def dataset_to_generator(
     dataset: Dataset,
@@ -84,3 +89,60 @@ def _get_number_of_logical_cores() -> int:
         )
         N = 0
     return N
+
+
+class MultiDataset(Dataset):
+    def __init__(self, datasets, transform=None):
+        """
+        Args:
+            datasets: A list of datasets to sample from.
+            transform: A function that takes N items (one from each dataset) and combines them.
+        """  # noqa: E501
+        self.datasets = datasets
+        self.transform = transform
+
+    def __len__(self):
+        # Length is defined by the smallest dataset in the list
+        return min(len(ds) for ds in self.datasets)
+
+    def __getitem__(self, idx):
+        sampled_items = [ds[idx] for ds in self.datasets]
+
+        if self.transform:
+            # Apply the transformation to all sampled items
+            return self.transform(*sampled_items)
+
+        return tuple(sampled_items)
+
+
+class ShuffledDataset(Dataset):
+    def __init__(self, dataset):
+        """
+        Wrapper that shuffles the dataset indices once.
+
+        Args:
+            dataset (Dataset): The original dataset to shuffle.
+        """
+        self.dataset = dataset
+        self.shuffled_indices = np.random.permutation(
+            len(dataset)
+        )  # Shuffle indices once
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        """
+        Returns the data at the shuffled index.
+
+        Args:
+            idx (int): Index in the shuffled dataset.
+        """
+        original_idx = self.shuffled_indices[idx]
+        return self.dataset[original_idx]
+
+
+def dataset_to_Xy(ds: Dataset):
+    return dataset_to_generator(ds, batch_size=len(ds), shuffle=False, num_workers=0)(
+        None
+    )

{imt_ring-1.6.31 → imt_ring-1.6.33}/tests/test_custom_joints.py

@@ -43,7 +43,7 @@ def test_virtual_input_joint_axes_rr_joint():
     np.testing.assert_allclose(
         -X["seg1"]["joint_axes"],
         np.repeat(np.array([[0.0, 1, 0]]), 1000, axis=0),
-        atol=4e-7,
+        atol=2e-6,
     )
     np.testing.assert_allclose(
         -X["seg3"]["joint_axes"],
@@ -57,8 +57,8 @@ def test_virtual_input_joint_axes_rr_joint():
     np.testing.assert_allclose(
         X["seg1"]["joint_axes"],
         np.repeat(-joint_axes[1:2], 1000, axis=0),
-        atol=4e-7,
-        rtol=2e-4,
+        atol=2e-6,
+        rtol=5e-4,
     )
     np.testing.assert_allclose(
         -X["seg3"]["joint_axes"],
@@ -103,7 +103,7 @@ def test_virtual_input_joint_axes_rr_imp_joint():
     np.testing.assert_allclose(
         -X["seg1"]["joint_axes"],
         np.repeat(np.array([[0.0, 1, 0]]), 1000, axis=0),
-        atol=4e-7,
+        atol=2e-6,
     )
     np.testing.assert_allclose(
         X["seg3"]["joint_axes"],

{imt_ring-1.6.31 → imt_ring-1.6.33}/tests/test_random.py

@@ -26,6 +26,8 @@ def test_delta_ang_min_max(
             dt, next_phi = _resolve_range_of_motion(
                 range_of_motion,
                 range_of_motion_method,
+                -2 * np.pi,
+                2 * np.pi,
                 0.0,
                 3.14,
                 float(delta_ang_min),
@@ -83,10 +85,10 @@ def test_angle(randomized_interpolation, range_of_motion, range_of_motion_method
 def test_position():
     for Ts in [0.1, 0.01]:
         T = 30
-        POS_0 = 0.0
-        pos = ring.algorithms.random_position_over_time(
-            jrand.PRNGKey(1), POS_0, -0.2, 0.2, 0.1, 0.5, 0.1, 0.5, T, Ts, None, 10
-        )
-        assert pos.shape == (int(T / Ts),)
-        # TODO Why does this fail for POS_0 != 0.0?
-        assert pos[0] == POS_0
+        for POS_0 in [0.0, 1.0]:
+            pos = ring.algorithms.random_position_over_time(
+                jrand.PRNGKey(1), POS_0, -0.2, 0.2, 0.1, 0.5, 0.1, 0.5, T, Ts, None, 10
+            )
+            assert pos.shape == (int(T / Ts),)
+            # TODO Why does this fail for POS_0 != 0.0?
+            assert pos[0] == POS_0

{imt_ring-1.6.31 → imt_ring-1.6.33}/tests/test_train.py

@@ -33,7 +33,8 @@ def test_rnno():
     ml.train_fn(
         gen,
         5,
-        ml.RNNO(N * 4, return_quats=True, eval=False, hidden_state_dim=20),
+        # .unwrapped to get ride of the `GroundtruthHeadingWrapper`
+        ml.RNNO(N * 4, return_quats=True, eval=False, hidden_state_dim=20).unwrapped,
     )