imt-ring 1.7.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

{imt_ring-1.7.0.dist-info → imt_ring-1.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: imt-ring
-Version: 1.7.0
+Version: 1.7.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.7.0.dist-info → imt_ring-1.7.1.dist-info}/RECORD

@@ -4,13 +4,13 @@ ring/base.py,sha256=AkG_Gpk7i2j77MzxzjiolJ9WNGcSq_3aOcpu0l6-0e0,50543
 ring/maths.py,sha256=R22SNQutkf9v7Hp9klo0wvJVIyBQz0O8_5oJaDQcFis,12652
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
-ring/algorithms/_random.py,sha256=UMyv-VPZLcErrKqs0XB83QJjs8GrmoNsv-zRSxGXvnI,14490
+ring/algorithms/_random.py,sha256=9mSP7M1On_CBz-dEbvE0iDzl_buf125EiYb1y10Ylxc,15805
 ring/algorithms/dynamics.py,sha256=NFOZawjrFoS5RgiWOpG1pQCC8l7RBOEZIi9ok6gvf9U,12268
-ring/algorithms/jcalc.py,sha256=l6BXOmXwrZ_AKKRm4gEHq_k2LSUQ4wd--1qL1qNTcKk,46794
+ring/algorithms/jcalc.py,sha256=zUMKEndSF6zaU7qHjJEU-iEQ_YAFNTNPcOwzdyu33uk,47632
 ring/algorithms/kinematics.py,sha256=IXeTQ-afzeEzLVmQVQ1oTXJxz_lTwvaWlgHeJkhO_8o,7423
 ring/algorithms/sensors.py,sha256=v_TZMyWjffDpPwoyS1fy8X-1i9y1vDf6mk1EmGS2ztc,18251
 ring/algorithms/custom_joints/__init__.py,sha256=3pQ-Is_HBTQDkzESCNg9VfoP8wvseWmooryG8ERnu_A,366
-ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwaMQyUjYv7EtsPiU3A,2402
+ring/algorithms/custom_joints/rr_imp_joint.py,sha256=1sOej-D3q4bSh5KJqcK5fZA3iOHXUlrlT9u-tIa-D6c,2351
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
 ring/algorithms/custom_joints/rsaddle_joint.py,sha256=QoMo6NXdYgA9JygSzBvr0eCdd3qKhUgCrGPNO2Qdxko,1200
 ring/algorithms/custom_joints/suntay.py,sha256=TZG307NqdMiXnNY63xEx8AkAjbQBQ4eO6DQ7R4j4D08,16726
@@ -28,7 +28,7 @@ ring/extras/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
 ring/extras/dataloader.py,sha256=dfNPjnxDoKxWGKSImuJ_49CWgBn73vxSEek8COq9nNk,3749
 ring/extras/dataloader_torch.py,sha256=t2DDiB9ZHb_SzFlVbntCGGIybj4F-NoA0PaB4_afjGw,3983
 ring/extras/hdf5.py,sha256=XPIrwogD-d544yy08UJyfLVp1ZKRUtiZukW7RA8VUxQ,5856
-ring/extras/interactive_viewer.py,sha256=vQEzcBDdG3BPqTGEktC74DsCfvgKktj9DKWK8gBzRtE,3805
+ring/extras/interactive_viewer.py,sha256=-jpoSsDrdzshZw-_MSI96QLvm9foRWTUS_aa_CNFk74,3867
 ring/extras/normalizer.py,sha256=o26stPP6EHasZQxQX0vKqTrhUNZBaJ2O17L6W_gBMN4,1699
 ring/extras/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
 ring/extras/torch_loss_fn.py,sha256=1LnWTmtxXPxoQFr4QixW12AjpRUfrseSDBmifhu6ErE,2676
@@ -89,8 +89,8 @@ ring/utils/__init__.py,sha256=Q37bjy2wjRGggd77MHlgl_50i2zOuVnPny4yOLiTe-8,567
 ring/utils/batchsize.py,sha256=uCj8LG7elbjEUUzuK29Z3I9T8bxJTcsybY3DdGeqhQs,1786
 ring/utils/path.py,sha256=zRPfxYNesvgefkddd26oar6f9433LkMGkhp9dF3rPUs,1926
 ring/utils/utils.py,sha256=gKwOXLxWraeZfX6EbBcg3hkq30DcXN0mcRUeOSTNiMo,7336
-imt_ring-1.7.0.dist-info/METADATA,sha256=CNwgvWr9Yu7MgIfcNwXkuByr7_8vkxvO5IkJg3iDKbs,5887
-imt_ring-1.7.0.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-imt_ring-1.7.0.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
-imt_ring-1.7.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.7.0.dist-info/RECORD,,
+imt_ring-1.7.1.dist-info/METADATA,sha256=cFujMyd1Xpqa8NiYn3iMqnhPHzGvlb_mLBbIPplQA6M,5887
+imt_ring-1.7.1.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+imt_ring-1.7.1.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
+imt_ring-1.7.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.7.1.dist-info/RECORD,,

ring/algorithms/_random.py

@@ -41,30 +41,48 @@ def random_angle_over_time(
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
     interpolation_method: str = "cosine",
+    include_standstills_prob: float = 0.0,  # 0.0 means no standstills
+    include_standstills_t_min: float = 0.5,
+    include_standstills_t_max: float = 5.0,
 ) -> jax.Array:
     def body_fn_outer(val):
         i, t, phi, key_t, key_ang, ANG = val
 
-        key_t, consume_t = random.split(key_t)
+        key_t, consume_t, consume_standstill = random.split(key_t, 3)
         key_ang, consume_ang = random.split(key_ang)
         rom_halfsize_float = _to_float(rom_halfsize, t)
         rom_lower = ANG_0 - rom_halfsize_float
         rom_upper = ANG_0 + rom_halfsize_float
-        dt, phi = _resolve_range_of_motion(
-            range_of_motion,
-            range_of_motion_method,
-            rom_lower,
-            rom_upper,
-            _to_float(dang_min, t),
-            _to_float(dang_max, t),
-            _to_float(delta_ang_min, t),
-            _to_float(delta_ang_max, t),
-            t_min,
-            _to_float(t_max, t),
-            phi,
-            consume_t,
-            consume_ang,
-            max_iter,
+
+        is_standstill = jax.random.bernoulli(
+            consume_standstill, include_standstills_prob
+        )
+        dt, phi = jax.lax.cond(
+            is_standstill,
+            lambda: (
+                jax.random.uniform(
+                    consume_t,
+                    minval=include_standstills_t_min,
+                    maxval=include_standstills_t_max,
+                ),
+                phi,
+            ),
+            lambda: _resolve_range_of_motion(
+                range_of_motion,
+                range_of_motion_method,
+                rom_lower,
+                rom_upper,
+                _to_float(dang_min, t),
+                _to_float(dang_max, t),
+                _to_float(delta_ang_min, t),
+                _to_float(delta_ang_max, t),
+                t_min,
+                _to_float(t_max, t),
+                phi,
+                consume_t,
+                consume_ang,
+                max_iter,
+            ),
         )
         t += dt
 
@@ -119,7 +137,8 @@ def random_angle_over_time(
 
 # APPROVED
 def random_position_over_time(
-    key: random.PRNGKey,
+    key_t: random.PRNGKey,
+    key_value: random.PRNGKey,
     POS_0: float,
     pos_min: float | TimeDependentFloat,
     pos_max: float | TimeDependentFloat,
@@ -135,19 +154,14 @@ def random_position_over_time(
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
     interpolation_method: str = "cosine",
+    include_standstills_prob: float = 0.0,  # 0.0 means no standstills
+    include_standstills_t_min: float = 0.5,
+    include_standstills_t_max: float = 5.0,
 ) -> jax.Array:
     def body_fn_inner(val):
         i, t, t_pre, x, x_pre, key = val
         dt = t - t_pre
 
-        def sample_dx_squared(key):
-            key, consume = random.split(key)
-            dx = (
-                random.uniform(consume) * (2 * dpos_max * t_max**2)
-                - dpos_max * t_max**2
-            )
-            return key, dx
-
         def sample_dx(key):
             key, consume1, consume2 = random.split(key, 3)
             sign = random.choice(consume1, jnp.array([-1.0, 1.0]))
@@ -182,24 +196,43 @@ def random_position_over_time(
         return jnp.logical_not(break_if_true1 | break_if_true2)
 
     def body_fn_outer(val):
-        i, t, t_pre, x, x_pre, key, POS = val
-        key, consume = random.split(key)
-        t += random.uniform(consume, minval=t_min, maxval=_to_float(t_max, t_pre))
+        i, t, t_pre, x, x_pre, key_t, key_value, POS = val
+        key_t, consume_t, consume_standstill = random.split(key_t, 3)
 
-        # that zero resets the max_it count
-        val_inner = (0, t, t_pre, x, x_pre, key)
-        _, t, t_pre, x, x_pre, key = jax.lax.while_loop(
-            cond_fn_inner, body_fn_inner, val_inner
+        is_standstill = jax.random.bernoulli(
+            consume_standstill, include_standstills_prob
+        )
+
+        def is_standstill_branch():
+            dt = random.uniform(
+                consume_t,
+                minval=include_standstills_t_min,
+                maxval=include_standstills_t_max,
+            )
+            t = t_pre + dt
+            return 0, t, t_pre, x, x_pre, key_value
+
+        def no_standstill_branch():
+            dt = random.uniform(consume_t, minval=t_min, maxval=_to_float(t_max, t_pre))
+            t = t_pre + dt
+            # that zero resets the max_it count
+            val_inner = (0, t, t_pre, x, x_pre, key_value)
+            return jax.lax.while_loop(cond_fn_inner, body_fn_inner, val_inner)
+
+        _, t, t_pre, x, x_pre, key_value = jax.lax.cond(
+            is_standstill,
+            is_standstill_branch,
+            no_standstill_branch,
         )
 
         POS_i = jnp.array([[jnp.floor(t / Ts) * Ts, x]])
         POS = jax.lax.dynamic_update_slice_in_dim(POS, POS_i, start_index=i, axis=0)
         t_pre = t
         x_pre = x
-        return i + 1, t, t_pre, x, x_pre, key, POS
+        return i + 1, t, t_pre, x, x_pre, key_t, key_value, POS
 
     def cond_fn_outer(val):
-        i, t, t_pre, x, x_pre, key, POS = val
+        i, t, t_pre, x, x_pre, key_t, key_value, POS = val
         return t <= T
 
     # preallocate POS array
@@ -207,7 +240,7 @@ def random_position_over_time(
     POS = jnp.zeros((int(T // t_min) + 1, 2))
     POS = POS.at[0, 1].set(POS_0)
 
-    val_outer = (1, 0.0, 0.0, POS_0, POS_0, key, POS)
+    val_outer = (1, 0.0, 0.0, POS_0, POS_0, key_t, key_value, POS)
     end, *_, consume, POS = jax.lax.while_loop(cond_fn_outer, body_fn_outer, val_outer)
     POS = jnp.where(
         (jnp.arange(len(POS)) < end)[:, None],

ring/algorithms/custom_joints/rr_imp_joint.py

@@ -23,11 +23,10 @@ def register_rr_imp_joint(
         return ring.Transform.create(rot=rot)
 
     def _draw_rr_imp(config, key_t, key_value, dt, N, _):
-        key_t1, key_t2 = jax.random.split(key_t)
         key_value1, key_value2 = jax.random.split(key_value)
-        q_traj_pri = _draw_rxyz(config, key_t1, key_value1, dt, N, _)
+        q_traj_pri = _draw_rxyz(config, key_t, key_value1, dt, N, _)
         q_traj_res = _draw_rxyz(
-            replace(config_res, T=config.T), key_t2, key_value2, dt, N, _
+            replace(config_res, T=config.T), key_t, key_value2, dt, N, _
         )
         # scale to be within bounds
         q_traj_res = q_traj_res * (jnp.deg2rad(ang_max_deg) / jnp.pi)

ring/algorithms/jcalc.py

@@ -174,6 +174,16 @@ class MotionConfig:
         default_factory=lambda: dict()
     )
 
+    # fields related to simulating standstills (no motion time periods)
+    # these are "Joint Standstills" so the standstills are calculated on
+    # a joint level, for each joint independently
+    # This means that a `standstills_prob` of 20% means that each joint
+    # has at each dt \in [t_min, t_max] drawing process a probability of
+    # 20% that it will just stay at its current joint value
+    include_standstills_prob: float = 0.0  # in %; 0% means no standstills
+    include_standstills_t_min: float = 0.5
+    include_standstills_t_max: float = 5.0
+
     def is_feasible(self) -> bool:
         return _is_feasible_config1(self)
 
@@ -791,12 +801,15 @@ def _draw_rxyz(
         config.cdf_bins_min,
         config.cdf_bins_max,
         config.interpolation_method,
+        config.include_standstills_prob,
+        config.include_standstills_t_min,
+        config.include_standstills_t_max,
     )
 
 
 def _draw_pxyz(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
@@ -811,6 +824,7 @@ def _draw_pxyz(
     )
     max_iter = 100
     return _random.random_position_over_time(
+        key_t,
         key_value,
         POS_0,
         config.cor_pos_min if cor else config.pos_min,
@@ -827,6 +841,9 @@ def _draw_pxyz(
         config.cdf_bins_min,
         config.cdf_bins_max,
         config.interpolation_method,
+        config.include_standstills_prob,
+        config.include_standstills_t_min,
+        config.include_standstills_t_max,
     )
 
 
@@ -840,7 +857,6 @@ def _draw_spherical(
 ) -> jax.Array:
     # NOTE: We draw 3 euler angles and then build a quaternion.
     # Not ideal, but i am unaware of a better way.
-    @jax.vmap
     def draw_euler_angles(key_t, key_value):
         return _draw_rxyz(
             config,
@@ -853,8 +869,9 @@ def _draw_spherical(
             free_spherical=True,
         )
 
-    triple = lambda key: jax.random.split(key, 3)
-    euler_angles = draw_euler_angles(triple(key_t), triple(key_value)).T
+    euler_angles = jax.vmap(draw_euler_angles, in_axes=(None, 0))(
+        key_t, jax.random.split(key_value, 3)
+    ).T
     q = maths.quat_euler(euler_angles)
     return q
 
@@ -867,7 +884,6 @@ def _draw_saddle(
     N: int | None,
     _: jax.Array,
 ) -> jax.Array:
-    @jax.vmap
     def draw_euler_angles(key_t, key_value):
         return _draw_rxyz(
             config,
@@ -880,14 +896,15 @@ def _draw_saddle(
             free_spherical=False,
         )
 
-    double = lambda key: jax.random.split(key)
-    yz_euler_angles = draw_euler_angles(double(key_t), double(key_value)).T
+    yz_euler_angles = jax.vmap(draw_euler_angles, in_axes=(None, 0))(
+        key_t, jax.random.split(key_value)
+    ).T
     return yz_euler_angles
 
 
 def _draw_p3d_and_cor(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
@@ -896,7 +913,7 @@ def _draw_p3d_and_cor(
 ) -> jax.Array:
     keys = jax.random.split(key_value, 3)
 
-    def draw(key, xyz: str):
+    def draw(key_value, xyz: str):
         return _draw_pxyz(
             replace(
                 config,
@@ -905,8 +922,8 @@ def _draw_p3d_and_cor(
                 pos0_min=getattr(config, f"pos0_min_p3d_{xyz}"),
                 pos0_max=getattr(config, f"pos0_max_p3d_{xyz}"),
             ),
-            None,
-            key,
+            key_t,
+            key_value,
             dt,
             N,
             None,
@@ -919,26 +936,26 @@ def _draw_p3d_and_cor(
 
 def _draw_p3d(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
     __: jax.Array,
 ) -> jax.Array:
-    return _draw_p3d_and_cor(config, _, key_value, dt, N, None, cor=False)
+    return _draw_p3d_and_cor(config, key_t, key_value, dt, N, None, cor=False)
 
 
 def _draw_cor(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
-    q_free = _draw_free(config, _, key_value1, dt, N, None)
-    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, N, None, cor=True)
+    q_free = _draw_free(config, key_t, key_value1, dt, N, None)
+    q_p3d = _draw_p3d_and_cor(config, key_t, key_value2, dt, N, None, cor=True)
     return jnp.concatenate((q_free, q_p3d), axis=1)
 
 
@@ -952,7 +969,7 @@ def _draw_free(
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
     q = _draw_spherical(config, key_t, key_value1, dt, N, None)
-    pos = _draw_p3d(config, None, key_value2, dt, N, None)
+    pos = _draw_p3d(config, key_t, key_value2, dt, N, None)
     return jnp.concatenate((q, pos), axis=1)
 
 
@@ -975,7 +992,7 @@ def _draw_free_2d(
         enable_range_of_motion=False,
         free_spherical=True,
     )[:, None]
-    pos_yz = _draw_p3d(config, None, key_value2, dt, N, None)[:, :2]
+    pos_yz = _draw_p3d(config, key_t, key_value2, dt, N, None)[:, :2]
     return jnp.concatenate((angle_x, pos_yz), axis=1)
 
 

ring/extras/interactive_viewer.py

@@ -82,10 +82,10 @@ def _fire_main(path_sys_xml: str, path_qs_np: Optional[str] = None, **scene_kwar
 
     assert qs.ndim == 2, f"qs.shape = {qs.shape}"
     T, Q = qs.shape
-    assert Q == sys.q_size()
+    assert Q == sys.q_size(), f"Q={Q} != sys.q_size={sys.q_size()}"
     dt_target = sys.dt
 
-    with InteractiveViewer(sys, **scene_kwargs) as viewer:
+    with InteractiveViewer(sys, width=640, height=480, **scene_kwargs) as viewer:
         dt = dt_target
         last_t = time.time()
         t = -1