imt-ring 1.3.5__tar.gz → 1.3.7__tar.gz

{imt_ring-1.3.5 → imt_ring-1.3.7}/PKG-INFO

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

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

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.5
+Version: 1.3.7
 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.3.5 → imt_ring-1.3.7}/src/ring/algorithms/custom_joints/suntay.py

@@ -293,12 +293,25 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
     ring.register_new_joint_type(name, joint_model, 1, overwrite=True)
 
 
+def _scale_delta(method: str, key, xs, mn, mx, amin, amax, **kwargs):
+    if method == "normal":
+        delta = jnp.clip(jax.random.normal(key) + 0.5, 1.0)
+    elif method == "uniform":
+        delta = 1 / (jax.random.uniform(key) + 1e-2)
+    else:
+        raise NotImplementedError
+
+    return delta
+
+
 def Polynomial_DrawnFnPair(
     order: int = 2,
     center: bool = False,
     flexion_center_deg: Optional[float] = None,
     include_bias: bool = True,
     enable_scale_delta: bool = True,
+    scale_delta_method: str = "normal",
+    scale_delta_kwargs: dict = dict(),
 ) -> DrawnFnPairFactory:
     assert not (order == 0 and not include_bias)
 
@@ -343,7 +356,9 @@ def Polynomial_DrawnFnPair(
             amin, amax = jnp.min(values), jnp.max(values) + eps
             if enable_scale_delta:
                 delta = amax - amin
-                scale_delta = jnp.clip(jax.random.normal(c3) + 0.5, 1.0)
+                scale_delta = _scale_delta(
+                    scale_delta_method, c3, xs, mn, mx, amin, amax, **scale_delta_kwargs
+                )
                 amax = amin + delta * scale_delta
             return amin, amax, poly_factors, q0
 

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/ring/algorithms/generator/base.py

@@ -140,8 +140,11 @@ class RCMG:
         batchsize: int = 1,
         sizes: int | list[int] = 1,
         seed: int = 1,
+        shuffle: bool = True,
     ) -> types.BatchedGenerator:
-        return batch.batch_generators_eager(self.gens, sizes, batchsize, seed=seed)
+        return batch.batch_generators_eager(
+            self.gens, sizes, batchsize, seed=seed, shuffle=shuffle
+        )
 
     def to_lazy_gen(
         self, sizes: int | list[int] = 1, jit: bool = True

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/ring/algorithms/sensors.py

@@ -3,6 +3,7 @@ from typing import Optional
 
 import jax
 import jax.numpy as jnp
+
 from ring import algebra
 from ring import algorithms
 from ring import base
@@ -445,22 +446,18 @@ def _joint_axes_from_xs(sys, xs, sys_xs):
 
 def _joint_axes_from_sys(sys: base.Transform, N: int) -> dict:
     "`sys` should be `sys_noimu`. `N` is number of timesteps"
-    xaxis = jnp.array([1.0, 0, 0])
-    yaxis = jnp.array([0.0, 1, 0])
-    zaxis = jnp.array([0.0, 0, 1])
-    id_to_axis = {"x": xaxis, "y": yaxis, "z": zaxis}
     X = {}
 
     def f(_, __, name, link_type, link):
         joint_params = link.joint_params
         if link_type in ["rx", "ry", "rz"]:
-            joint_axes = id_to_axis[link_type[1]]
+            joint_axes = maths.unit_vectors(link_type[1])
         elif link_type == "rr":
             joint_axes = joint_params["rr"]["joint_axes"]
         elif link_type[:6] == "rr_imp":
             joint_axes = joint_params[link_type]["joint_axes"]
         else:
-            joint_axes = xaxis
+            joint_axes = maths.x_unit_vector
         X[name] = {"joint_axes": joint_axes}
 
     sys.scan(f, "lll", sys.link_names, sys.link_types, sys.links)

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/ring/base.py

@@ -571,6 +571,7 @@ class System(_Base):
         new_damp: Optional[jax.Array] = None,
         new_stif: Optional[jax.Array] = None,
         new_zero: Optional[jax.Array] = None,
+        seed: int = 1,
     ):
         "By default damping, stiffness are set to zero."
         from ring.algorithms import get_joint_model
@@ -600,7 +601,7 @@ class System(_Base):
 
         jm = get_joint_model(new_joint_type)
         if jm.init_joint_params is not None:
-            sys = sys.from_str(sys.to_str())
+            sys = sys.from_str(sys.to_str(), seed=seed)
 
         return sys
 

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/ring/ml/train.py

@@ -5,6 +5,8 @@ from typing import Callable, Optional, Tuple
 import jax
 import jax.numpy as jnp
 import optax
+import tree_utils
+
 from ring import maths
 from ring.algorithms.generator import types
 from ring.ml import base as ml_base
@@ -15,8 +17,6 @@ from ring.utils import distribute_batchsize
 from ring.utils import expand_batchsize
 from ring.utils import parse_path
 from ring.utils import pickle_load
-import tree_utils
-
 import wandb
 
 # (T, N, F) -> Scalar
@@ -142,15 +142,17 @@ def train_fn(
         Wether or not the training run was killed by a callback.
     """
 
+    filter = filter.nojit()
+
     if checkpoint is not None:
         checkpoint = Path(checkpoint).with_suffix(".pickle")
         recv_checkpoint: dict = pickle_load(checkpoint)
-        filter.params = recv_checkpoint["params"]
+        filter_params = recv_checkpoint["params"]
         opt_state = recv_checkpoint["opt_state"]
+        del recv_checkpoint
+    else:
+        filter_params = filter.search_attr("params")
 
-    filter = filter.nojit()
-
-    filter_params = filter.search_attr("params")
     if filter_params is None:
         X, _ = generator(jax.random.PRNGKey(1))
         filter_params, _ = filter.init(X=X, seed=seed_network)

{imt_ring-1.3.5 → imt_ring-1.3.7}/src/ring/utils/utils.py

@@ -122,6 +122,8 @@ def pytree_deepcopy(tree):
         return tuple(pytree_deepcopy(ele) for ele in tree)
     elif isinstance(tree, dict):
         return {key: pytree_deepcopy(value) for key, value in tree.items()}
+    elif isinstance(tree, _Base):
+        return tree
     else:
         raise NotImplementedError(f"Not implemented for type={type(tree)}")
 

{imt_ring-1.3.5 → imt_ring-1.3.7}/tests/test_custom_joints.py

@@ -1,5 +1,7 @@
 import jax
 import numpy as np
+import pytest
+
 import ring
 from ring import MotionConfig
 from ring.algorithms import custom_joints
@@ -21,14 +23,19 @@ def pipeline_load_data_X(
     return gen(jax.random.PRNGKey(1))[0]
 
 
-def test_virtual_input_joint_axes_rr_joint():
-    sys = ring.io.load_example("test_three_seg_seg2")
+def _replace_ry_rz_with(sys, new_joint_type: str):
     sys_rr = sys.replace(
         link_types=[
-            "rr" if link_type in ["ry", "rz"] else link_type
+            new_joint_type if link_type in ["ry", "rz"] else link_type
             for link_type in sys.link_types
         ]
     )
+    return sys_rr
+
+
+def test_virtual_input_joint_axes_rr_joint():
+    sys = ring.io.load_example("test_three_seg_seg2")
+    sys_rr = _replace_ry_rz_with(sys, "rr")
     sys_rr = _init_joint_params(jax.random.PRNGKey(1), sys_rr)
     joint_axes = sys_rr.links.joint_params["rr"]["joint_axes"]
 
@@ -38,7 +45,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=1e-7,
+        atol=4e-7,
     )
     np.testing.assert_allclose(
         X["seg3"]["joint_axes"],
@@ -52,8 +59,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=3e-7,
-        rtol=2e-6,
+        atol=4e-7,
+        rtol=2e-4,
     )
     np.testing.assert_allclose(
         X["seg3"]["joint_axes"],
@@ -63,6 +70,7 @@ def test_virtual_input_joint_axes_rr_joint():
     )
 
 
+@pytest.mark.filterwarnings("ignore:The system has")
 def test_virtual_input_joint_axes_rr_imp_joint():
     custom_joints.register_rr_imp_joint(MotionConfig(T=10.0))
 
@@ -85,8 +93,8 @@ def test_virtual_input_joint_axes_rr_imp_joint():
     np.testing.assert_allclose(
         X["seg3"]["joint_axes"],
         np.repeat(-joint_axes[3:4], 1000, axis=0),
-        atol=0.002,
-        rtol=0.002,
+        atol=0.01,
+        rtol=0.032,
     )
 
     # test `make_generator`
@@ -97,10 +105,10 @@ 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=1e-7,
+        atol=4e-7,
     )
     np.testing.assert_allclose(
         X["seg3"]["joint_axes"],
         np.repeat(np.array([[0.0, 0, 1]]), 1000, axis=0),
-        atol=1e-7,
+        atol=4e-7,
     )

{imt_ring-1.3.5 → imt_ring-1.3.7}/tests/test_random.py

@@ -1,6 +1,7 @@
 from jax import random as jrand
 import numpy as np
 import pytest
+
 import ring
 from ring.algorithms._random import _resolve_range_of_motion
 
@@ -17,21 +18,27 @@ from ring.algorithms._random import _resolve_range_of_motion
 def test_delta_ang_min_max(
     range_of_motion, range_of_motion_method, delta_ang_min, delta_ang_max
 ):
+    max_iter = 100
+    t_min, t_max = 0.1, 3.1
     for seed in range(10):
-        key = jrand.PRNGKey(seed)
-        phi = _resolve_range_of_motion(
-            range_of_motion,
-            range_of_motion_method,
-            0.0,
-            3.14,
-            float(delta_ang_min),
-            float(delta_ang_max),
-            1.0,
-            0.0,
-            key,
-            100,
-        )
-        assert delta_ang_min < abs(phi) < delta_ang_max
+        for prev_phi in [-3.0, 0.0, 3.0]:
+            key_t, key_ang = jrand.split(jrand.PRNGKey(seed))
+            dt, next_phi = _resolve_range_of_motion(
+                range_of_motion,
+                range_of_motion_method,
+                0.0,
+                3.14,
+                float(delta_ang_min),
+                float(delta_ang_max),
+                t_min,
+                t_max,
+                prev_phi,
+                key_t,
+                key_ang,
+                max_iter,
+            )
+            assert t_min <= dt <= t_max
+            assert delta_ang_min < abs(next_phi - prev_phi) < delta_ang_max
 
 
 @pytest.mark.parametrize(