imt-ring 1.5.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

{imt_ring-1.5.0.dist-info → imt_ring-1.5.1.dist-info}/METADATA

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

@@ -14,10 +14,10 @@ ring/algorithms/custom_joints/rr_imp_joint.py,sha256=a3JT0w7pB94kZ95eBR8ZO853eSe
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
 ring/algorithms/custom_joints/suntay.py,sha256=7-kym1kMDwqYD_2um1roGcBeB8BlTCPe1wljuNGNARA,16676
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=1rzClXZ0WMJ_IZroTO9i1aWiHBy7whsCrJLIMY4zC3c,13280
+ring/algorithms/generator/base.py,sha256=LOckKDehObDwgOSu_uBhlqkxdztJ0NHTP2mbxwtEcwQ,13335
 ring/algorithms/generator/batch.py,sha256=Hwh5jYZQEmkx73YaXjWd6sZdikmj43spE7DCzGDHXtE,6637
 ring/algorithms/generator/finalize_fns.py,sha256=0fbtwQw89_w0ytQ_aJ877CZGY5fbtb8sbsRO0O8pT34,9081
-ring/algorithms/generator/motion_artifacts.py,sha256=vzBLlG60KCAa7Zj1RdUiRkoOx_3inA_2M1mBKl3lTKs,8834
+ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
 ring/algorithms/generator/setup_fns.py,sha256=MFz3czHBeWs1Zk1A8O02CyQpQ-NCyW9PMpbqmKit6es,1455
 ring/algorithms/generator/types.py,sha256=HjNyATFSLfHkXlzdJhvUkiqnhzpXFDDXmWS3LYBlOtU,721
@@ -81,7 +81,7 @@ ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
 ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
 ring/utils/utils.py,sha256=VkB0Gvmlaz2MZdntgjWA0rOpRkvIRpLWRFgIofoY7hs,5441
-imt_ring-1.5.0.dist-info/METADATA,sha256=ZjJMt4357zV4eK-ZKH_d4Q7nhc8dJ6RG_AnocvzDNzU,3104
-imt_ring-1.5.0.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-imt_ring-1.5.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.5.0.dist-info/RECORD,,
+imt_ring-1.5.1.dist-info/METADATA,sha256=C4QFyeh8L1nslQdrwJ6tBtL_dqDRGsUl7W8ZExZA2hc,3104
+imt_ring-1.5.1.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+imt_ring-1.5.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.5.1.dist-info/RECORD,,

ring/algorithms/generator/base.py

@@ -34,7 +34,7 @@ class RCMG:
         randomize_motion_artifacts: bool = False,
         randomize_joint_params: bool = False,
         imu_motion_artifacts: bool = False,
-        imu_motion_artifacts_kwargs: dict = dict(hide_injected_bodies=True),
+        imu_motion_artifacts_kwargs: dict = dict(),
         dynamic_simulation: bool = False,
         dynamic_simulation_kwargs: dict = dict(),
         output_transform: Optional[Callable] = None,
@@ -50,9 +50,6 @@ class RCMG:
         for c in config:
             assert c.is_feasible()
 
-        if cor:
-            sys = [s._replace_free_with_cor() for s in sys]
-
         self.gens = []
         for _sys in sys:
             self.gens.append(
@@ -78,6 +75,7 @@ class RCMG:
                     output_transform=output_transform,
                     keep_output_extras=keep_output_extras,
                     use_link_number_in_Xy=use_link_number_in_Xy,
+                    cor=cor,
                 )
             )
 
@@ -238,6 +236,7 @@ def _build_mconfig_batched_generator(
     output_transform: Callable | None,
     keep_output_extras: bool,
     use_link_number_in_Xy: bool,
+    cor: bool,
 ) -> types.BatchedGenerator:
 
     if add_X_jointaxes or add_y_relpose or add_y_rootincl:
@@ -284,13 +283,17 @@ def _build_mconfig_batched_generator(
         for f in pipe:
             key, consume = jax.random.split(key)
             sys = f(consume, sys)
+        if cor:
+            sys = sys._replace_free_with_cor()
         return sys
 
     def _finalize_fn(Xy: types.Xy, extras: types.OutputExtras):
         pipe = []
         if dynamic_simulation:
             pipe.append(finalize_fns.DynamicalSimulation(**dynamic_simulation_kwargs))
-        if imu_motion_artifacts and imu_motion_artifacts_kwargs["hide_injected_bodies"]:
+        if imu_motion_artifacts and imu_motion_artifacts_kwargs.get(
+            "hide_injected_bodies", True
+        ):
             pipe.append(motion_artifacts.HideInjectedBodies())
         if finalize_fn is not None:
             pipe.append(finalize_fns.FinalizeFn(finalize_fn))
@@ -312,15 +315,15 @@ def _build_mconfig_batched_generator(
         return Xy, extras
 
     def _gen(key: types.PRNGKey):
+        key, *consume = jax.random.split(key, len(config) + 1)
+        syss = jax.vmap(_setup_fn, (0, None))(jnp.array(consume), sys)
+
         qs = []
-        for _config in config:
-            key, _q = draw_random_q(key, sys, _config)
+        for i, _config in enumerate(config):
+            key, _q = draw_random_q(key, syss[i], _config)
             qs.append(_q)
         qs = jnp.stack(qs)
 
-        key, *consume = jax.random.split(key, len(config) + 1)
-        syss = jax.vmap(_setup_fn, (0, None))(jnp.array(consume), sys)
-
         @jax.vmap
         def _vmapped_context(key, q, sys):
             x, _ = jax.vmap(kinematics.forward_kinematics_transforms, (None, 0))(sys, q)

ring/algorithms/generator/motion_artifacts.py

@@ -1,3 +1,4 @@
+import inspect
 import warnings
 
 import jax
@@ -127,6 +128,7 @@ def setup_fn_randomize_damping_stiffness_factory(
     prob_rigid: float = 0.0,
     all_imus_either_rigid_or_flex: bool = False,
     imus_surely_rigid: list[str] = [],
+    **kwargs,
 ):
     assert 0 <= prob_rigid <= 1
     assert prob_rigid != 1, "Use `imu_motion_artifacts`=False instead."
@@ -198,6 +200,18 @@ def setup_fn_randomize_damping_stiffness_factory(
     return setup_fn_randomize_damping_stiffness
 
 
+# assert that there exists no keyword arg duplicate which would induce ambiguity
+kwargs = lambda f: set(inspect.signature(f).parameters.keys())
+assert (
+    len(
+        kwargs(inject_subsystems).intersection(
+            kwargs(setup_fn_randomize_damping_stiffness_factory)
+        )
+    )
+    == 1
+)
+
+
 def _match_q_x_between_sys(
     sys_small: base.System,
     q_large: jax.Array,