imt-ring 1.3.3__py3-none-any.whl → 1.3.11__py3-none-any.whl

{imt_ring-1.3.3.dist-info → imt_ring-1.3.11.dist-info}/METADATA

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

@@ -1,23 +1,23 @@
 ring/__init__.py,sha256=iNvbAZi7Qfa69IbL1z4lB7zHL8WusV5fBrKah2la-Gc,1566
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=ZAoe9B1HbAX9NYiKaisssTBn-1VBXoJTsWgFAvlQoZw,33705
+ring/base.py,sha256=YFPrUWelWswEhq8x8Byv-5pK64mipiGW6x5IlMr4we4,33803
 ring/maths.py,sha256=jJr_kr78-XDce8B4tXQ2Li-jBntVQhaS8csxglCsj8A,12193
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=t3YXcgqMJxadUjFiILVD0HlQRPLdrQyc8aKiB36w0vE,1701
-ring/algorithms/_random.py,sha256=6EG0GHYe6tCq0qUt4Jes8W1EaqqaLa0sSZhnwBbEjCE,13340
-ring/algorithms/dynamics.py,sha256=nqq5I0RYSbHNlGiLMlohz08IfL9Njsrid4upDnwkGbI,10629
-ring/algorithms/jcalc.py,sha256=oqSiwz3Be1VfIpmJXEFTNM_9_o3tyuTtyZt2aqttyN4,28213
+ring/algorithms/_random.py,sha256=M9JQSMXSUARWuzlRLP3Wmkuntrk9LZpP30p4_IPgDB4,13805
+ring/algorithms/dynamics.py,sha256=_TwclBXe6vi5C5iJWAIeUIJEIMHQ_1QTmnHvCEpVO0M,10867
+ring/algorithms/jcalc.py,sha256=6olMYQtgKZE5KBEAHF0Rqxe__1wcZQVEiLgm1vO7_Gw,28260
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
-ring/algorithms/sensors.py,sha256=Y3Wo9qj3BWKoIHB0V04QwyD-Z5m4BrAjfBX8Pn6y9Lg,18005
+ring/algorithms/sensors.py,sha256=MICO9Sn0AfoqRx_9KWR3hufsIID-K6SOIg3oPDgsYMU,17869
 ring/algorithms/custom_joints/__init__.py,sha256=fzeE7TdUhmGgbbFAyis1tKcyQ4Fo8LigDwD3hUVnH_w,316
 ring/algorithms/custom_joints/rr_imp_joint.py,sha256=a3JT0w7pB94kZ95eBR8ZO853eSeyjFoiXmhYlaXoHDE,2392
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
-ring/algorithms/custom_joints/suntay.py,sha256=0Oym3KQssj3QDOldnz9PTy5jPg9ZLk85mMK2YX1qvB4,15600
+ring/algorithms/custom_joints/suntay.py,sha256=7-kym1kMDwqYD_2um1roGcBeB8BlTCPe1wljuNGNARA,16676
 ring/algorithms/generator/__init__.py,sha256=p4ucl0zQtp5NwNoXIRjmTzGGRu2WOAWFfNmYRPwQles,912
-ring/algorithms/generator/base.py,sha256=zmrRK_I6BWoo4WbEcEVK7iFKdPfetc6txs7U8iu1xEk,14771
-ring/algorithms/generator/batch.py,sha256=BGzmwH1AItXjPRyHtsYnAfYnoogw8jxhng9oyVw72lw,9019
-ring/algorithms/generator/motion_artifacts.py,sha256=aKdkZU5OF4_aKyL4Yo-ftZRwrDCve1LuuREGAUlTqtI,8551
-ring/algorithms/generator/pd_control.py,sha256=3pOaYig26vmp8gippDfy2KNJRZO3kr0rGd_PBIuEROM,5759
+ring/algorithms/generator/base.py,sha256=sr-YZkjd8pZJAI5vFG_IqOO4AEeiEYtXr8uUsPMS6Q4,14779
+ring/algorithms/generator/batch.py,sha256=bslFSN2Gs_aX9cNwFooExhKUwevc70q3bspEMTwygm4,9256
+ring/algorithms/generator/motion_artifacts.py,sha256=_kiAl1VHoX1fW5AUlXOtPBWyHIIFof_M78AP-m9f1ME,8790
+ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
 ring/algorithms/generator/randomize.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
 ring/algorithms/generator/transforms.py,sha256=nvNDvM20tEw9Zd0ra0TxA25uf01L40Y2UKvtQmOrlGo,12782
 ring/algorithms/generator/types.py,sha256=CAhvDK5qiHnrGtkCVccB07doiz_D6lHJ35B7sW0pyZA,1110
@@ -50,13 +50,14 @@ ring/io/xml/from_xml.py,sha256=8b44sPVWgoY8JGJZLpJ8M_eLfcfu3IsMtBzSytPTPmw,9234
 ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,1562
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
 ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
-ring/ml/__init__.py,sha256=-bryExVoKJYSF_G_KYc5hI_GciIhj2xZ8WGi6TdRghw,1836
+ring/ml/__init__.py,sha256=52LpEjni5lG-ov5-3ocodH-vKZxNcFMU7W9XfjDicp0,2113
 ring/ml/base.py,sha256=PQ72VasEqlecBZgWP5HE5rWYyLiLq7nCVLymXo9f0dw,8959
-ring/ml/callbacks.py,sha256=DkSy5c7IRqAAks2dx8acEBExYxUv-xiUFwZn4odPYq4,13253
+ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
-ring/ml/ringnet.py,sha256=OWRDu2COmptzbpJWlRLbPIn_ioKZCAd_iu-eiY_aPjk,8521
-ring/ml/train.py,sha256=ftt2MOSSNGCdL7ZoAXcbIgeHW1Wkpgp6XYyLIBUIClI,10872
+ring/ml/ringnet.py,sha256=rgje5AKUKpT8K-vbE9_SgZ3IijR8TJEHnaqxsE57Mhc,8617
+ring/ml/rnno_v1.py,sha256=T4SKG7iypqn2HBQLKhDmJ2Slj2Z5jtUBHvX_6aL8pyM,1103
+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
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
@@ -70,14 +71,15 @@ ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
 ring/sys_composer/inject_sys.py,sha256=Mj-q-mUjXKwkg-ol6IQAjf9IJfk7pGhez0_WoTKTgm0,3503
 ring/sys_composer/morph_sys.py,sha256=2GpPtS5hT0eZMptdGpt30Hc97OykJNE67lEVRf7sHrc,12700
-ring/utils/__init__.py,sha256=rTvSA4RiJAVCY_A64FUMd8IJTv94LgoSA3Ps5X63_jA,799
-ring/utils/batchsize.py,sha256=mPFGD7AedFMycHtyIuZtNWCaAvKLLWSWaB7X6u54xvM,1358
+ring/utils/__init__.py,sha256=FZ9ziQrWlx16QIpQ8RdLKrvN_17CAdvnZMNNodxWY0o,812
+ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
+ring/utils/batchsize.py,sha256=ByXGX7bw2gwrVirEsazm2JXnwNPGnqgEirzziYoSUS0,1553
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
-ring/utils/utils.py,sha256=I2f6-DMBrrgy5tpLzPLlezifQgkO2fERZWyX3cfb4sI,5303
-imt_ring-1.3.3.dist-info/METADATA,sha256=nTihurycKYmLCI61Cojd7VLrnb1gpd-H8nwUupysaC8,3104
-imt_ring-1.3.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.3.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.3.dist-info/RECORD,,
+ring/utils/utils.py,sha256=mIcKNv5v2de8HrG7bAhl2bNfmwkMZyIIwFkJq2XWMOI,5357
+imt_ring-1.3.11.dist-info/METADATA,sha256=5-bAEaCMi5JiwNMOEshFJlV1ISU3840hImimdj7l2CU,3105
+imt_ring-1.3.11.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.3.11.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.3.11.dist-info/RECORD,,

ring/algorithms/_random.py

@@ -4,6 +4,7 @@ import warnings
 import jax
 from jax import random
 import jax.numpy as jnp
+
 from ring import maths
 
 Float = jax.Array
@@ -40,20 +41,20 @@ def random_angle_over_time(
     def body_fn_outer(val):
         i, t, phi, key_t, key_ang, ANG = val
 
-        key_t, consume = random.split(key_t)
-        dt = random.uniform(consume, minval=t_min, maxval=_to_float(t_max, t))
-
-        key_ang, consume = random.split(key_ang)
-        phi = _resolve_range_of_motion(
+        key_t, consume_t = random.split(key_t)
+        key_ang, consume_ang = random.split(key_ang)
+        dt, phi = _resolve_range_of_motion(
             range_of_motion,
             range_of_motion_method,
             _to_float(dang_min, t),
             _to_float(dang_max, t),
             _to_float(delta_ang_min, t),
             _to_float(delta_ang_max, t),
-            dt,
+            t_min,
+            _to_float(t_max, t),
             phi,
-            consume,
+            consume_t,
+            consume_ang,
             max_iter,
         )
         t += dt
@@ -246,12 +247,14 @@ def _resolve_range_of_motion(
     dang_max,
     delta_ang_min,
     delta_ang_max,
-    dt,
+    t_min,
+    t_max,
     prev_phi,
-    key,
+    key_t,
+    key_ang,
     max_iter,
 ):
-    def _next_phi(key):
+    def _next_phi(key, dt):
         key, consume = random.split(key)
 
         if range_of_motion:
@@ -294,21 +297,33 @@ def _resolve_range_of_motion(
             return prev_phi + sign * dphi
 
     def body_fn(val):
-        key, _, i = val
-        key, consume = jax.random.split(key)
-        next_phi = _next_phi(consume)
-        return key, next_phi, i + 1
+        key_t, key_ang, _, _, i = val
+
+        key_t, consume_t = jax.random.split(key_t)
+        dt = jax.random.uniform(consume_t, minval=t_min, maxval=t_max)
+
+        key_ang, consume_ang = jax.random.split(key_ang)
+        next_phi = _next_phi(consume_ang, dt)
+
+        return key_t, key_ang, dt, next_phi, i + 1
 
     def cond_fn(val):
-        _, next_phi, i = val
+        *_, dt, next_phi, i = val
         delta_phi = jnp.abs(next_phi - prev_phi)
-        # delta is in bounds
-        break_if_true1 = (delta_phi >= delta_ang_min) & (delta_phi <= delta_ang_max)
+        # delta_ang is in bounds
+        cond_delta_ang = (delta_phi >= delta_ang_min) & (delta_phi <= delta_ang_max)
+        # dang is in bounds
+        dang = delta_phi / dt
+        cond_dang = (dang >= dang_min) & (dang <= dang_max)
+
+        break_if_true1 = jnp.logical_and(cond_delta_ang, cond_dang)
+        # break out of loop
         break_if_true2 = i > max_iter
         return (i == 0) | (jnp.logical_not(break_if_true1 | break_if_true2))
 
-    # the `prev_phi` here is unused
-    return jax.lax.while_loop(cond_fn, body_fn, (key, prev_phi, 0))[1]
+    init_val = (key_t, key_ang, 1.0, prev_phi, 0)
+    *_, dt, next_phi, _ = jax.lax.while_loop(cond_fn, body_fn, init_val)
+    return dt, next_phi
 
 
 def cosInterpolate(x, xp, fp):

ring/algorithms/custom_joints/suntay.py

@@ -293,16 +293,36 @@ 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,
-    val: float = 2.0,
     center: bool = False,
-    flexion_center: Optional[float] = None,
+    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 val >= 0.0
+    assert not (order == 0 and not include_bias)
+
+    flexion_center = (
+        jnp.deg2rad(flexion_center_deg) if flexion_center_deg is not None else None
+    )
+    del flexion_center_deg
 
     # because 0-th order is also counted
     order += 1
+    powers = jnp.arange(order) if include_bias else jnp.arange(1, order)
 
     def factory(xs, mn, mx):
         nonlocal flexion_center
@@ -311,17 +331,22 @@ def Polynomial_DrawnFnPair(
         flexion_mx = jnp.max(xs)
 
         def _apply_poly_factors(poly_factors, q):
-            return poly_factors @ jnp.power(q, jnp.arange(order))
+            return poly_factors @ jnp.power(q, powers)
 
         if flexion_center is None:
             flexion_center = (flexion_mn + flexion_mx) / 2
+
+        if (order - 1) == 0:
+            method = "clip"
+            minval, maxval = mn, mx
         else:
-            flexion_center = jnp.array(flexion_center)
+            method = "minmax"
+            minval, maxval = -1.0, 1.0
 
         def init(key):
-            c1, c2 = jax.random.split(key)
+            c1, c2, c3 = jax.random.split(key, 3)
             poly_factors = jax.random.uniform(
-                c1, shape=(order,), minval=-val, maxval=val
+                c1, shape=(len(powers),), minval=minval, maxval=maxval
             )
             q0 = jax.random.uniform(c2, minval=flexion_mn, maxval=flexion_mx)
             values = jax.vmap(_apply_poly_factors, in_axes=(None, 0))(
@@ -329,13 +354,19 @@ def Polynomial_DrawnFnPair(
             )
             eps = 1e-6
             amin, amax = jnp.min(values), jnp.max(values) + eps
+            if enable_scale_delta:
+                delta = amax - amin
+                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
 
         def _apply(params, q):
             amin, amax, poly_factors, q0 = params
             q = q - q0
             value = _apply_poly_factors(poly_factors, q)
-            return restrict(value, mn, mx, amin, amax)
+            return restrict(value, mn, mx, amin, amax, method=method)
 
         if center:
 

ring/algorithms/dynamics.py

@@ -1,7 +1,9 @@
 from typing import Optional, Tuple
+import warnings
 
 import jax
 import jax.numpy as jnp
+
 from ring import algebra
 from ring import base
 from ring import maths
@@ -213,7 +215,7 @@ def forward_dynamics(
     q: jax.Array,
     qd: jax.Array,
     tau: jax.Array,
-    mass_mat_inv: jax.Array,
+    # mass_mat_inv: jax.Array,
 ) -> Tuple[jax.Array, jax.Array]:
     C = inverse_dynamics(sys, qd, jnp.zeros_like(qd))
     mass_matrix = compute_mass_matrix(sys)
@@ -235,6 +237,11 @@ def forward_dynamics(
 
         mass_mat_inv = jax.scipy.linalg.solve(mass_matrix, eye, assume_a="pos")
     else:
+        warnings.warn(
+            f"You are using `sys.mass_mat_iters`={sys.mass_mat_iters} which is >0. "
+            "This feature is currently not fully supported. See the local TODO."
+        )
+        mass_mat_inv = jnp.diag(jnp.ones((sys.qd_size(),)))
         mass_mat_inv = _inv_approximate(mass_matrix, mass_mat_inv, sys.mass_mat_iters)
 
     return mass_mat_inv @ qf_smooth, mass_mat_inv
@@ -254,9 +261,8 @@ def _strapdown_integration(
 def _semi_implicit_euler_integration(
     sys: base.System, state: base.State, taus: jax.Array
 ) -> base.State:
-    qdd, mass_mat_inv = forward_dynamics(
-        sys, state.q, state.qd, taus, state.mass_mat_inv
-    )
+    qdd, mass_mat_inv = forward_dynamics(sys, state.q, state.qd, taus)
+    del mass_mat_inv
     qd_next = state.qd + sys.dt * qdd
 
     q_next = []
@@ -277,7 +283,7 @@ def _semi_implicit_euler_integration(
     sys.scan(q_integrate, "qdl", state.q, qd_next, sys.link_types)
     q_next = jnp.concatenate(q_next)
 
-    state = state.replace(q=q_next, qd=qd_next, mass_mat_inv=mass_mat_inv)
+    state = state.replace(q=q_next, qd=qd_next)
     return state
 
 

ring/algorithms/generator/base.py

@@ -4,6 +4,7 @@ import warnings
 
 import jax
 import jax.numpy as jnp
+import tqdm
 import tree_utils
 
 from ring import base
@@ -83,10 +84,14 @@ class RCMG:
             ), "If `randomize_anchors`, then only one system is expected"
             sys = randomize.randomize_anchors(sys[0], **randomize_anchors_kwargs)
 
-        zip_sys_config = False
         if randomize_hz:
-            zip_sys_config = True
             sys, config = randomize.randomize_hz(sys, config, **randomize_hz_kwargs)
+        else:
+            # create zip
+            N_sys = len(sys)
+            sys = sum([len(config) * [s] for s in sys], start=[])
+            config = N_sys * config
+            assert len(sys) == len(config)
 
         if sys_ml is None:
             # TODO
@@ -97,17 +102,10 @@ class RCMG:
             sys_ml = sys[0]
 
         self.gens = []
-        if zip_sys_config:
-            for _sys, _config in zip(sys, config):
-                self.gens.append(
-                    partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
-                )
-        else:
-            for _sys in sys:
-                for _config in config:
-                    self.gens.append(
-                        partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
-                    )
+        for _sys, _config in tqdm.tqdm(
+            zip(sys, config), desc="building generators", total=len(sys)
+        ):
+            self.gens.append(partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml))
 
     def _to_data(self, sizes, seed):
         return batch.batch_generators_eager_to_list(self.gens, sizes, seed=seed)
@@ -140,8 +138,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

ring/algorithms/generator/batch.py

@@ -63,12 +63,12 @@ def batch_generators_lazy(
 
 
 def _number_of_executions_required(size: int) -> int:
-    vmap_threshold = 128
     _, vmap = utils.distribute_batchsize(size)
 
+    eager_threshold = utils.batchsize_thresholds()[1]
     primes = iter(utils.primes(vmap))
     n_calls = 1
-    while vmap > vmap_threshold:
+    while vmap > eager_threshold:
         prime = next(primes)
         n_calls *= prime
         vmap /= prime
@@ -86,7 +86,11 @@ def batch_generators_eager_to_list(
 
     key = jax.random.PRNGKey(seed)
     data = []
-    for gen, size in tqdm(zip(generators, sizes), desc="eager data generation"):
+    for gen, size in tqdm(
+        zip(generators, sizes),
+        desc="executing generators",
+        total=len(sizes),
+    ):
 
         n_calls = _number_of_executions_required(size)
         # decrease size by n_calls times
@@ -97,8 +101,10 @@ def batch_generators_eager_to_list(
         for _ in range(n_calls):
             key, consume = jax.random.split(key)
             sample = gen_jit(consume)
-            # converts also to numpy
+            # converts also to numpy; but with np.array.flags.writeable = False
             sample = jax.device_get(sample)
+            # this then sets this flag to True
+            sample = jax.tree_map(np.array, sample)
             data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
 
     return data
@@ -145,7 +151,7 @@ def _data_fn_from_paths(
     paths = [utils.parse_path(p, mkdir=False) for p in paths]
 
     extensions = list(set([Path(p).suffix for p in paths]))
-    assert len(extensions) == 1
+    assert len(extensions) == 1, f"{extensions}"
 
     if extensions[0] == ".h5":
         N = sum([utils.hdf5_load_length(p) for p in paths])

ring/algorithms/generator/motion_artifacts.py

@@ -49,6 +49,7 @@ def inject_subsystems(
     rotational_damp: float = 0.1,
     translational_stif: float = 50.0,
     translational_damp: float = 0.1,
+    disable_warning: bool = False,
     **kwargs,
 ) -> base.System:
     imu_idx_to_name_map = {sys.name_to_idx(imu): imu for imu in sys.findall_imus()}
@@ -92,10 +93,11 @@ def inject_subsystems(
     # TODO set all joint_params to zeros; they can not be preserved anyways and
     # otherwise many warnings will be rose
     # instead warn explicitly once now and move on
-    warnings.warn(
-        "`sys.links.joint_params` has been set to zero, this might lead to "
-        "unexpected behaviour unless you use `randomize_joint_params`"
-    )
+    if not disable_warning:
+        warnings.warn(
+            "`sys.links.joint_params` has been set to zero, this might lead to "
+            "unexpected behaviour unless you use `randomize_joint_params`"
+        )
     joint_params_zeros = tree_utils.tree_zeros_like(sys.links.joint_params)
     sys = sys.replace(links=sys.links.replace(joint_params=joint_params_zeros))
 
@@ -180,9 +182,13 @@ def setup_fn_randomize_damping_stiffness_factory(
             link_spring_stiffness = link_spring_stiffness.at[slice].set(stif)
             link_damping = link_damping.at[slice].set(damp)
 
-        assert len(imus_surely_rigid) == len(triggered_surely_rigid)
+        assert len(imus_surely_rigid) == len(
+            triggered_surely_rigid
+        ), f"{imus_surely_rigid}, {triggered_surely_rigid}"
         for imu_surely_rigid in imus_surely_rigid:
-            assert imu_surely_rigid in triggered_surely_rigid
+            assert (
+                imu_surely_rigid in triggered_surely_rigid
+            ), f"{imus_surely_rigid} not in {triggered_surely_rigid}"
 
         return sys.replace(
             link_damping=link_damping, link_spring_stiffness=link_spring_stiffness

ring/algorithms/generator/pd_control.py

@@ -4,6 +4,7 @@ from typing import Optional
 from flax import struct
 import jax
 import jax.numpy as jnp
+
 from ring import base
 from ring.algorithms import dynamics
 from ring.algorithms import jcalc
@@ -49,7 +50,7 @@ def _pd_control(P: jax.Array, D: Optional[jax.Array] = None):
         assert sys.q_size() == q_ref.shape[1], f"q_ref.shape = {q_ref.shape}"
         assert sys.qd_size() == P.size
         if D is not None:
-            sys.qd_size() == D.size
+            assert sys.qd_size() == D.size
 
         q_ref_as_dict = {}
         qd_ref_as_dict = {}

ring/algorithms/jcalc.py

@@ -424,14 +424,14 @@ def _draw_rxyz(
     # TODO, delete these args and pass a modifified `config` with `replace` instead
     enable_range_of_motion: bool = True,
     free_spherical: bool = False,
+    # how often it should try to fullfill the dang_min/max and delta_ang_min/max conds
+    max_iter: int = 5,
 ) -> jax.Array:
     key_value, consume = jax.random.split(key_value)
     ANG_0 = jax.random.uniform(consume, minval=config.ang0_min, maxval=config.ang0_max)
     # `random_angle_over_time` always returns wrapped angles, thus it would be
     # inconsistent to allow an initial value that is not wrapped
     ANG_0 = maths.wrap_to_pi(ANG_0)
-    # only used for `delta_ang_min_max` logic
-    max_iter = 5
     return _random.random_angle_over_time(
         key_t,
         key_value,

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)

ring/base.py

@@ -571,8 +571,11 @@ 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
+
         q_size, qd_size = Q_WIDTHS[new_joint_type], QD_WIDTHS[new_joint_type]
 
         def logic_unfreeze_to_spherical(link_name, olt, ola, old, ols, olz):
@@ -594,7 +597,13 @@ class System(_Base):
 
             return nlt, nla, nld, nls, nlz
 
-        return _update_sys_if_replace_joint_type(self, logic_unfreeze_to_spherical)
+        sys = _update_sys_if_replace_joint_type(self, logic_unfreeze_to_spherical)
+
+        jm = get_joint_model(new_joint_type)
+        if jm.init_joint_params is not None:
+            sys = sys.from_str(sys.to_str(), seed=seed)
+
+        return sys
 
     def findall_imus(self) -> list[str]:
         return [name for name in self.link_names if name[:3] == "imu"]
@@ -988,13 +997,11 @@ class State(_Base):
         q (jax.Array): System state in minimal coordinates (equals `sys.q_size()`)
         qd (jax.Array): System velocity in minimal coordinates (equals `sys.qd_size()`)
         x: (Transform): Maximal coordinates of all links. From epsilon-to-link.
-        mass_mat_inv (jax.Array): Inverse of the mass matrix. Internal usage.
     """
 
     q: jax.Array
     qd: jax.Array
     x: Transform
-    mass_mat_inv: jax.Array
 
     @classmethod
     def create(
@@ -1048,4 +1055,4 @@ class State(_Base):
         if x is None:
             x = Transform.zero((sys.num_links(),))
 
-        return cls(q, qd, x, jnp.diag(jnp.ones((sys.qd_size(),))))
+        return cls(q, qd, x)

ring/ml/__init__.py

@@ -3,6 +3,7 @@ from . import callbacks
 from . import ml_utils
 from . import optimizer
 from . import ringnet
+from . import rnno_v1
 from . import train
 from . import training_loop
 from .base import AbstractFilter
@@ -42,17 +43,28 @@ def RNNO(
     params=None,
     eval: bool = True,
     samp_freq: float | None = None,
+    v1: bool = False,
     **kwargs,
 ):
     assert "message_dim" not in kwargs
     assert "link_output_normalize" not in kwargs
     assert "link_output_dim" not in kwargs
 
+    if v1:
+        kwargs.update(
+            dict(forward_factory=rnno_v1.rnno_v1_forward_factory, output_dim=output_dim)
+        )
+    else:
+        kwargs.update(
+            dict(
+                message_dim=0,
+                link_output_normalize=False,
+                link_output_dim=output_dim,
+            )
+        )
+
     ringnet = RING(  # noqa: F811
         params=params,
-        message_dim=0,
-        link_output_normalize=False,
-        link_output_dim=output_dim,
         **kwargs,
     )
     ringnet = base.NoGraph_FilterWrapper(ringnet, quat_normalize=return_quats)

ring/ml/callbacks.py

@@ -245,7 +245,8 @@ class SaveParamsTrainingLoopCallback(training_loop.TrainingLoopCallback):
             else:
                 value = "{:.2f}".format(ele.value).replace(".", ",")
                 filename = parse_path(
-                    self.path_to_file + f"_episode={ele.episode}_value={value}",
+                    str(Path(self.path_to_file).with_suffix(""))
+                    + f"_episode={ele.episode}_value={value}",
                     extension="pickle",
                 )
 
@@ -404,7 +405,7 @@ class CheckpointCallback(training_loop.TrainingLoopCallback):
         # only checkpoint if run has been killed
         if training_loop.recv_kill_run_signal():
             path = parse_path(
-                "~/.xxy_checkpoints", ml_utils.unique_id(), extension="pickle"
+                "~/.ring_checkpoints", ml_utils.unique_id(), extension="pickle"
             )
             data = {"params": self.params, "opt_state": self.opt_state}
             pickle_save(

ring/ml/ringnet.py

@@ -191,8 +191,16 @@ class LSTM(hk.RNNCore):
 
 
 class RING(ml_base.AbstractFilter):
-    def __init__(self, params=None, lam=None, jit: bool = True, name=None, **kwargs):
-        self.forward_lam_factory = partial(make_ring, **kwargs)
+    def __init__(
+        self,
+        params=None,
+        lam=None,
+        jit: bool = True,
+        name=None,
+        forward_factory=make_ring,
+        **kwargs,
+    ):
+        self.forward_lam_factory = partial(forward_factory, **kwargs)
         self.params = self._load_params(params)
         self.lam = lam
         self._name = name

ring/ml/rnno_v1.py

@@ -0,0 +1,41 @@
+from typing import Optional, Sequence
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+def rnno_v1_forward_factory(
+    output_dim: int,
+    rnn_layers: Sequence[int] = (400, 300),
+    linear_layers: Sequence[int] = (200, 100, 50, 50, 25, 25),
+    layernorm: bool = True,
+    act_fn_linear=jax.nn.relu,
+    act_fn_rnn=jax.nn.elu,
+    lam: Optional[tuple[int]] = None,
+):
+    # unused
+    del lam
+
+    @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):
+            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)
+            hk.set_state(f"rnn_{i}", state)
+
+            if layernorm:
+                X = hk.LayerNorm(axis=-1, create_scale=False, create_offset=False)(X)
+            X = act_fn_rnn(X)
+
+        for n_units in linear_layers:
+            X = hk.Linear(n_units)(X)
+            X = act_fn_linear(X)
+
+        y = hk.Linear(output_dim)(X)
+        return y[..., None, :]
+
+    return forward_fn

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
@@ -13,10 +15,7 @@ from ring.ml import ml_utils
 from ring.ml import training_loop
 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 +141,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)
@@ -215,7 +216,7 @@ def train_fn(
 
         callbacks_all.append(
             ml_callbacks.SaveParamsTrainingLoopCallback(
-                path_to_file=parse_path(callback_save_params, extension=""),
+                path_to_file=callback_save_params,
                 last_n_params=3,
                 track_metrices=callback_save_params_track_metrices,
                 cleanup=False,

ring/utils/__init__.py

@@ -1,4 +1,4 @@
-from .batchsize import backend
+from .batchsize import batchsize_thresholds
 from .batchsize import distribute_batchsize
 from .batchsize import expand_batchsize
 from .batchsize import merge_batchsize

ring/utils/backend.py

@@ -0,0 +1,30 @@
+import os
+import re
+
+
+def set_host_device_count(n):
+    """
+    By default, XLA considers all CPU cores as one device. This utility tells XLA
+    that there are `n` host (CPU) devices available to use. As a consequence, this
+    allows parallel mapping in JAX :func:`jax.pmap` to work in CPU platform.
+
+    .. note:: This utility only takes effect at the beginning of your program.
+        Under the hood, this sets the environment variable
+        `XLA_FLAGS=--xla_force_host_platform_device_count=[num_devices]`, where
+        `[num_device]` is the desired number of CPU devices `n`.
+
+    .. warning:: Our understanding of the side effects of using the
+        `xla_force_host_platform_device_count` flag in XLA is incomplete. If you
+        observe some strange phenomenon when using this utility, please let us
+        know through our issue or forum page. More information is available in this
+        `JAX issue <https://github.com/google/jax/issues/1408>`_.
+
+    :param int n: number of CPU devices to use.
+    """
+    xla_flags = os.getenv("XLA_FLAGS", "")
+    xla_flags = re.sub(
+        r"--xla_force_host_platform_device_count=\S+", "", xla_flags
+    ).split()
+    os.environ["XLA_FLAGS"] = " ".join(
+        ["--xla_force_host_platform_device_count={}".format(n)] + xla_flags
+    )

ring/utils/batchsize.py

@@ -1,19 +1,37 @@
-from typing import Optional, Tuple
+from typing import Tuple, TypeVar
 
 import jax
-from tree_utils import PyTree
+
+PyTree = TypeVar("PyTree")
+
+
+def batchsize_thresholds():
+    backend = jax.default_backend()
+    if backend == "cpu":
+        vmap_size_min = 1
+        eager_threshold = 4
+    elif backend == "gpu":
+        vmap_size_min = 8
+        eager_threshold = 32
+    else:
+        raise Exception(
+            f"Backend {backend} has no default values, please add them in this function"
+        )
+    return vmap_size_min, eager_threshold
 
 
 def distribute_batchsize(batchsize: int) -> Tuple[int, int]:
     """Distributes batchsize accross pmap and vmap."""
-    vmap_size_min = 8
+    vmap_size_min = batchsize_thresholds()[0]
     if batchsize <= vmap_size_min:
         return 1, batchsize
     else:
         n_devices = jax.local_device_count()
-        assert (
-            batchsize % n_devices
-        ) == 0, f"Your GPU count of {n_devices} does not split batchsize {batchsize}"
+        msg = (
+            f"Your local device count of {n_devices} does not split batchsize"
+            + f" {batchsize}. local devices are {jax.local_devices()}"
+        )
+        assert (batchsize % n_devices) == 0, msg
         vmap_size = int(batchsize / n_devices)
         return int(batchsize / vmap_size), vmap_size
 
@@ -35,17 +53,3 @@ def expand_batchsize(tree: PyTree, pmap_size: int, vmap_size: int) -> PyTree:
         ),
         tree,
     )
-
-
-CPU_ONLY = False
-
-
-def backend(cpu_only: bool = False, n_gpus: Optional[int] = None):
-    "Sets backend for all jax operations (including this library)."
-    global CPU_ONLY
-
-    if cpu_only and not CPU_ONLY:
-        CPU_ONLY = True
-        from jax import config
-
-        config.update("jax_platform_name", "cpu")

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)}")