imt-ring 1.6.16__tar.gz → 1.6.18__tar.gz

{imt_ring-1.6.16 → imt_ring-1.6.18}/PKG-INFO

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

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

{imt_ring-1.6.16 → imt_ring-1.6.18}/src/imt_ring.egg-info/PKG-INFO

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

@@ -20,6 +20,7 @@ src/ring/algorithms/sensors.py
 src/ring/algorithms/custom_joints/__init__.py
 src/ring/algorithms/custom_joints/rr_imp_joint.py
 src/ring/algorithms/custom_joints/rr_joint.py
+src/ring/algorithms/custom_joints/rsaddle_joint.py
 src/ring/algorithms/custom_joints/suntay.py
 src/ring/algorithms/generator/__init__.py
 src/ring/algorithms/generator/base.py
@@ -85,6 +86,7 @@ src/ring/utils/backend.py
 src/ring/utils/batchsize.py
 src/ring/utils/colab.py
 src/ring/utils/dataloader.py
+src/ring/utils/dataloader_torch.py
 src/ring/utils/hdf5.py
 src/ring/utils/normalizer.py
 src/ring/utils/path.py

{imt_ring-1.6.16 → imt_ring-1.6.18}/src/ring/__init__.py

@@ -121,6 +121,7 @@ _UNIQUE_ID = None
 def setup(
     rr_joint_kwargs: None | dict = dict(),
     rr_imp_joint_kwargs: None | dict = dict(),
+    rsaddle_joint_kwargs: None | dict = dict(),
     suntay_joint_kwargs: None | dict = None,
     train_timing_start: None | float = None,
     unique_id: None | str = None,
@@ -138,6 +139,9 @@ def setup(
     if rr_imp_joint_kwargs is not None:
         custom_joints.register_rr_imp_joint(**rr_imp_joint_kwargs)
 
+    if rsaddle_joint_kwargs is not None:
+        custom_joints.register_rsaddle_joint(**rsaddle_joint_kwargs)
+
     if suntay_joint_kwargs is not None:
         custom_joints.register_suntay(**suntay_joint_kwargs)
 

{imt_ring-1.6.16 → imt_ring-1.6.18}/src/ring/algorithms/custom_joints/__init__.py

@@ -1,5 +1,6 @@
 from .rr_imp_joint import register_rr_imp_joint
 from .rr_joint import register_rr_joint
+from .rsaddle_joint import register_rsaddle_joint
 from .suntay import ConstantValue_DrawnFnPair
 from .suntay import GP_DrawFnPair
 from .suntay import MLP_DrawnFnPair

imt_ring-1.6.18/src/ring/algorithms/custom_joints/rsaddle_joint.py

@@ -0,0 +1,40 @@
+import jax.numpy as jnp
+
+import ring
+from ring import maths
+from ring.algorithms.jcalc import _draw_saddle
+from ring.algorithms.jcalc import _p_control_term_rxyz
+from ring.algorithms.jcalc import _qd_from_q_cartesian
+
+
+def register_rsaddle_joint():
+    def _transform(q, params):
+        axes = params["joint_axes"]
+        rot1 = maths.quat_rot_axis(axes[0], q[0])
+        rot2 = maths.quat_rot_axis(axes[1], q[1])
+        rot = maths.quat_mul(rot2, rot1)
+        return ring.Transform.create(rot=rot)
+
+    def _motion_fn_gen(i: int):
+        def _motion_fn(params):
+            axis = params["joint_axes"][i]
+            return ring.base.Motion.create(ang=axis)
+
+        return _motion_fn
+
+    joint_model = ring.JointModel(
+        _transform,
+        motion=[_motion_fn_gen(i) for i in range(2)],
+        rcmg_draw_fn=_draw_saddle,
+        p_control_term=_p_control_term_rxyz,
+        qd_from_q=_qd_from_q_cartesian,
+        init_joint_params=_draw_random_joint_axes,
+    )
+
+    ring.register_new_joint_type("rsaddle", joint_model, 2, overwrite=True)
+
+
+def _draw_random_joint_axes(key):
+    return dict(
+        joint_axes=maths.rotate(jnp.array([1.0, 0, 0]), maths.quat_random(key, (2,)))
+    )

{imt_ring-1.6.16 → imt_ring-1.6.18}/src/ring/algorithms/generator/finalize_fns.py

@@ -160,6 +160,7 @@ _P_gains = {
     "spherical": jnp.array(3 * [P_rot]),
     "p3d": jnp.array(3 * [P_pos]),
     "saddle": jnp.array([P_rot, P_rot]),
+    "rsaddle": jnp.array([P_rot, P_rot]),
     "frozen": jnp.array([]),
     "suntay": jnp.array([P_rot]),
 }

{imt_ring-1.6.16 → imt_ring-1.6.18}/src/ring/algorithms/jcalc.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from dataclasses import asdict
 from dataclasses import dataclass
 from dataclasses import field
@@ -87,11 +88,39 @@ class MotionConfig:
         assert nomotion_config.is_feasible()
         return nomotion_config
 
+    @staticmethod
+    def overwrite_for_joint_type(joint_type: str, **changes) -> None:
+        """Changes values of the `MotionConfig` used by the draw_fn for only a specific
+        joint.
+        """
+        previous_changes = _overwrite_for_joint_type_changes[joint_type]
+        for change in changes:
+            assert change not in previous_changes, f"For jointtype={joint_type} you "
+            f"previously changed the value={change}. You can't change it again, this "
+            "is not supported."
+        previous_changes.update(changes)
+
+        jm = get_joint_model(joint_type)
+
+        def draw_fn(config, *args):
+            return jm.rcmg_draw_fn(replace(config, **changes), *args)
+
+        register_new_joint_type(
+            joint_type,
+            replace(jm, rcmg_draw_fn=draw_fn),
+            base.Q_WIDTHS[joint_type],
+            base.QD_WIDTHS[joint_type],
+            overwrite=True,
+        )
+
     @staticmethod
     def from_register(name: str) -> "MotionConfig":
         return _registered_motion_configs[name]
 
 
+_overwrite_for_joint_type_changes: dict[str, dict] = defaultdict(lambda: dict())
+
+
 _registered_motion_configs = {
     "hinUndHer": MotionConfig(
         t_min=0.3,

imt_ring-1.6.18/src/ring/utils/dataloader_torch.py

@@ -0,0 +1,62 @@
+import os
+
+import jax
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from tree_utils import PyTree
+
+from ring.utils import parse_path
+from ring.utils import pickle_load
+
+
+class FolderOfPickleFilesDataset(Dataset):
+    def __init__(self, path, transform=None):
+        self.files = self.listdir(path)
+        self.transform = transform
+        self.N = len(self.files)
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, idx: int):
+        element = pickle_load(self.files[idx])
+        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 dataset_to_generator(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle=True,
+    seed: int = 1,
+    **kwargs,
+):
+    torch.manual_seed(seed)
+
+    dl = DataLoader(
+        dataset,
+        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