imt-ring 1.6.21__tar.gz → 1.6.23__tar.gz

{imt_ring-1.6.21 → imt_ring-1.6.23}/PKG-INFO

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

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

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.21
+Version: 1.6.23
 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.21 → imt_ring-1.6.23}/src/ring/algorithms/_random.py

@@ -35,6 +35,9 @@ def random_angle_over_time(
     randomized_interpolation: bool = False,
     range_of_motion: bool = False,
     range_of_motion_method: str = "uniform",
+    # this value has nothing to do with `range_of_motion` flag
+    # this forces the value to stay within [ANG_0 - rom_halfsize, ANG_0 + rom_halfsize]
+    rom_halfsize: float | TimeDependentFloat = 2 * jnp.pi,
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
     interpolation_method: str = "cosine",
@@ -44,9 +47,14 @@ def random_angle_over_time(
 
         key_t, consume_t = random.split(key_t)
         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),
@@ -251,6 +259,8 @@ def _clip_to_pi(phi):
 def _resolve_range_of_motion(
     range_of_motion,
     range_of_motion_method,
+    rom_lower: float,
+    rom_upper: float,
     dang_min,
     dang_max,
     delta_ang_min,
@@ -265,44 +275,47 @@ def _resolve_range_of_motion(
     def _next_phi(key, dt):
         key, consume = random.split(key)
 
-        if range_of_motion:
-            if range_of_motion_method == "coinflip":
-                probs = jnp.array([0.5, 0.5])
-            elif range_of_motion_method == "uniform":
-                p = 0.5 * (1 - prev_phi / jnp.pi)
-                probs = jnp.array([p, (1 - p)])
-            elif range_of_motion_method[:7] == "sigmoid":
-                scale = 1.5
-                provided_params = range_of_motion_method.split("-")
-                if len(provided_params) == 2:
-                    scale = float(provided_params[-1])
-                hardcut = jnp.pi - 0.01
-                p = jnp.where(
-                    prev_phi > hardcut,
-                    0.0,
-                    jnp.where(
-                        prev_phi < -hardcut, 1.0, jax.nn.sigmoid(-scale * prev_phi)
-                    ),
-                )
-                probs = jnp.array([p, (1 - p)])
-            else:
-                raise NotImplementedError
+        # legacy reasons, without range of motion the `sign` value, so going
+        # left or right is 50-50 for free joints and spherical joints
+        if not range_of_motion:
+            range_of_motion_method = "coinflip"
+
+        if range_of_motion_method == "coinflip":
+            probs = jnp.array([0.5, 0.5])
+        elif range_of_motion_method == "uniform":
+            p = 0.5 * (1 - prev_phi / jnp.pi)
+            probs = jnp.array([p, (1 - p)])
+        elif range_of_motion_method[:7] == "sigmoid":
+            scale = 1.5
+            provided_params = range_of_motion_method.split("-")
+            if len(provided_params) == 2:
+                scale = float(provided_params[-1])
+            hardcut = jnp.pi - 0.01
+            p = jnp.where(
+                prev_phi > hardcut,
+                0.0,
+                jnp.where(prev_phi < -hardcut, 1.0, jax.nn.sigmoid(-scale * prev_phi)),
+            )
+            probs = jnp.array([p, (1 - p)])
+        else:
+            raise NotImplementedError
 
-            sign = random.choice(consume, jnp.array([1.0, -1.0]), p=probs)
-            lower = _clip_to_pi(prev_phi + sign * dang_min * dt)
-            upper = _clip_to_pi(prev_phi + sign * dang_max * dt)
+        sign = random.choice(consume, jnp.array([1.0, -1.0]), p=probs)
+        lower = prev_phi + sign * dang_min * dt
+        upper = prev_phi + sign * dang_max * dt
 
-            # swap if lower > upper
-            lower, upper = jnp.sort(jnp.hstack((lower, upper)))
+        if range_of_motion:
+            lower, upper = _clip_to_pi(lower), _clip_to_pi(upper)
 
-            key, consume = random.split(key)
-            return random.uniform(consume, minval=lower, maxval=upper)
+        # swap if lower > upper
+        lower, upper = jnp.sort(jnp.hstack((lower, upper)))
 
-        else:
-            dphi = random.uniform(consume, minval=dang_min, maxval=dang_max) * dt
-            key, consume = random.split(key)
-            sign = random.choice(consume, jnp.array([1.0, -1.0]))
-            return prev_phi + sign * dphi
+        # clip bounds given by the angular velocity bounds to the rom bounds
+        lower = jnp.clip(lower, a_min=rom_lower)
+        upper = jnp.clip(upper, a_max=rom_upper)
+
+        key, consume = random.split(key)
+        return random.uniform(consume, minval=lower, maxval=upper)
 
     def body_fn(val):
         key_t, key_ang, _, _, i = val

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/ring/algorithms/dynamics.py

@@ -190,11 +190,11 @@ def _spring_force(sys: base.System, q: jax.Array):
 
     def _calc_spring_force_per_link(_, __, q, zeropoint, typ):
         # cor is (free, p3d) stacked; free is (spherical, p3d) stacked
-        if typ in ["free", "cor"]:
+        if base.System.joint_type_is_free_or_cor(typ):
             quat_force = _quaternion_spring_force(zeropoint[:4], q[:4])
             pos_force = zeropoint[4:] - q[4:]
             q_spring_force_link = jnp.concatenate((quat_force, pos_force))
-        elif typ == "spherical":
+        elif base.System.joint_type_is_spherical(typ):
             q_spring_force_link = _quaternion_spring_force(zeropoint, q)
         else:
             q_spring_force_link = zeropoint - q
@@ -268,11 +268,11 @@ def _semi_implicit_euler_integration(
     q_next = []
 
     def q_integrate(_, __, q, qd, typ):
-        if typ in ["free", "cor"]:
+        if sys.joint_type_is_free_or_cor(typ):
             quat_next = _strapdown_integration(q[:4], qd[:3], sys.dt)
             pos_next = q[4:] + qd[3:] * sys.dt
             q_next_i = jnp.concatenate((quat_next, pos_next))
-        elif typ == "spherical":
+        elif sys.joint_type_is_spherical(typ):
             quat_next = _strapdown_integration(q, qd, sys.dt)
             q_next_i = quat_next
         else:

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/ring/algorithms/generator/base.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from functools import partial
 import random
 from typing import Callable, Optional
@@ -446,7 +447,15 @@ def draw_random_q(
         draw_fn = jcalc.get_joint_model(link_type).rcmg_draw_fn
         if draw_fn is None:
             raise Exception(f"The joint type {link_type} has no draw fn specified.")
-        q_link = draw_fn(config, key_t, key_value, sys.dt, N, joint_params)
+
+        if link_type in config.joint_type_specific_overwrites:
+            _config = replace(
+                config, **config.joint_type_specific_overwrites[link_type]
+            )
+        else:
+            _config = config
+
+        q_link = draw_fn(_config, key_t, key_value, sys.dt, N, joint_params)
         # even revolute and prismatic joints must be 2d arrays
         q_link = q_link if q_link.ndim == 2 else q_link[:, None]
         q_list.append(q_link)

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/ring/algorithms/jcalc.py

@@ -40,6 +40,12 @@ class MotionConfig:
     dpos_max: float | TimeDependentFloat = 0.7
     pos_min: float | TimeDependentFloat = -2.5
     pos_max: float | TimeDependentFloat = +2.5
+    pos_min_p3d_x: float | TimeDependentFloat = -2.5
+    pos_max_p3d_x: float | TimeDependentFloat = +2.5
+    pos_min_p3d_y: float | TimeDependentFloat = -2.5
+    pos_max_p3d_y: float | TimeDependentFloat = +2.5
+    pos_min_p3d_z: float | TimeDependentFloat = -2.5
+    pos_max_p3d_z: float | TimeDependentFloat = +2.5
 
     # used by both `random_angle_*` and `random_pos_*`
     # only used if `randomized_interpolation` is set
@@ -54,11 +60,22 @@ class MotionConfig:
     range_of_motion_hinge: bool = True
     range_of_motion_hinge_method: str = "uniform"
 
+    # this value has nothing to do with `range_of_motion` flag
+    # this forces the value to stay within [ANG_0 - rom_halfsize, ANG_0 + rom_halfsize]
+    # used only by the `_draw_rxyz` function
+    rom_halfsize: float | TimeDependentFloat = 2 * jnp.pi
+
     # initial value of joints
     ang0_min: float = -jnp.pi
     ang0_max: float = jnp.pi
     pos0_min: float = 0.0
     pos0_max: float = 0.0
+    pos0_min_p3d_x: float = 0.0
+    pos0_max_p3d_x: float = 0.0
+    pos0_min_p3d_y: float = 0.0
+    pos0_max_p3d_y: float = 0.0
+    pos0_min_p3d_z: float = 0.0
+    pos0_max_p3d_z: float = 0.0
 
     # cor (center of rotation) custom fields
     cor_t_min: float = 0.2
@@ -67,6 +84,14 @@ class MotionConfig:
     cor_dpos_max: float | TimeDependentFloat = 0.5
     cor_pos_min: float | TimeDependentFloat = -0.4
     cor_pos_max: float | TimeDependentFloat = 0.4
+    cor_pos0_min: float = 0.0
+    cor_pos0_max: float = 0.0
+
+    # specify changes for this motionconfig and for specific joint types
+    # map of `link_types` -> dictionary of changes
+    joint_type_specific_overwrites: dict[str, dict[str, Any]] = field(
+        default_factory=lambda: dict()
+    )
 
     def is_feasible(self) -> bool:
         return _is_feasible_config1(self)
@@ -92,6 +117,9 @@ class MotionConfig:
     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.
+        !!! Note
+            This applies these changes to *all* MotionConfigs for this joint type!
+            This takes precedence *over* `Motionconfig.joint_type_specific_overwrites`!
         """
         previous_changes = _overwrite_for_joint_type_changes[joint_type]
         for change in changes:
@@ -113,6 +141,56 @@ class MotionConfig:
             overwrite=True,
         )
 
+    @staticmethod
+    def overwrite_for_subsystem(
+        sys: base.System, link_name: str, **changes
+    ) -> base.System:
+        """Modifies motionconfig of all joints in subsystem with root `link_name`.
+        Note that if the subsystem contains a free joint then the jointtype will
+        will be re-named to `free_<link_name>`, then the RCMG flag `cor` will
+        potentially not work as expected because it searches for all joints of
+        type `free` to replace with `cor`. The workaround here is to change the
+        type already from `free` to `cor in the xml file.
+        This takes precedence *over* `Motionconfig.joint_type_specific_overwrites`!
+
+        Args:
+            sys (base.System): System object that gets updated
+            link_name (str): Root node of subsystem
+            changes: Changes to apply to the motionconfig
+
+        Return:
+            base.System: Updated system with new jointtypes
+        """
+        from ring.algorithms.generator.finalize_fns import _P_gains
+
+        # all bodies in the subsystem
+        bodies = sys.findall_bodies_subsystem(link_name) + [sys.name_to_idx(link_name)]
+
+        jts_subsys = set([sys.link_types[i] for i in bodies]) - set(["frozen"])
+        postfix = "_" + link_name
+        # create new joint types with updated motionconfig
+        for typ in jts_subsys:
+            register_new_joint_type(
+                typ + postfix,
+                get_joint_model(typ),
+                base.Q_WIDTHS[typ],
+                base.QD_WIDTHS[typ],
+            )
+            MotionConfig.overwrite_for_joint_type(typ + postfix, **changes)
+            _P_gains[typ + postfix] = _P_gains[typ]
+
+        # rename all jointtypes
+        new_link_types = [
+            (
+                sys.link_types[i] + postfix
+                if (i in bodies and sys.link_types[i] != "frozen")
+                else sys.link_types[i]
+            )
+            for i in range(sys.num_links())
+        ]
+        sys = sys.replace(link_types=new_link_types)
+        return sys
+
     @staticmethod
     def from_register(name: str) -> "MotionConfig":
         return _registered_motion_configs[name]
@@ -221,6 +299,37 @@ _registered_motion_configs = {
 }
 
 
+def _joint_specific_overwrites_free_cor(
+    id: str, dang: float, dpos: float
+) -> MotionConfig:
+    changes = dict(
+        dang_max_free_spherical=dang,
+        dpos_max=dpos,
+        cor_dpos_max=dpos,
+        t_min=1.5,
+        t_max=15.0,
+    )
+    return replace(
+        _registered_motion_configs[id],
+        joint_type_specific_overwrites=dict(free=changes, cor=changes),
+    )
+
+
+_registered_motion_configs.update(
+    {
+        f"{id}-S": _joint_specific_overwrites_free_cor(id, 0.2, 0.1)
+        for id in ["expSlow", "expFast", "hinUndHer", "standard"]
+    }
+)
+_registered_motion_configs.update(
+    {
+        f"{id}-S+": _joint_specific_overwrites_free_cor(id, 0.1, 0.05)
+        for id in ["expSlow", "expFast", "hinUndHer", "standard"]
+    }
+)
+del _joint_specific_overwrites_free_cor
+
+
 def _is_feasible_config1(c: MotionConfig) -> bool:
     t_min, t_max = c.t_min, _to_float(c.t_max, 0.0)
 
@@ -254,8 +363,29 @@ def _is_feasible_config1(c: MotionConfig) -> bool:
     cond2 = inside_box_checks(
         _to_float(c.pos_min, 0.0), _to_float(c.pos_max, 0.0), c.pos0_min, c.pos0_max
     )
+    cond3 = inside_box_checks(
+        _to_float(c.pos_min_p3d_x, 0.0),
+        _to_float(c.pos_max_p3d_x, 0.0),
+        c.pos0_min_p3d_x,
+        c.pos0_max_p3d_x,
+    )
+    cond4 = inside_box_checks(
+        _to_float(c.pos_min_p3d_y, 0.0),
+        _to_float(c.pos_max_p3d_y, 0.0),
+        c.pos0_min_p3d_y,
+        c.pos0_max_p3d_y,
+    )
+    cond5 = inside_box_checks(
+        _to_float(c.pos_min_p3d_z, 0.0),
+        _to_float(c.pos_max_p3d_z, 0.0),
+        c.pos0_min_p3d_z,
+        c.pos0_max_p3d_z,
+    )
+
+    # test that the delta_ang_min is smaller than 2*rom_halfsize
+    cond6 = _to_float(c.delta_ang_min, 0.0) < 2 * _to_float(c.rom_halfsize, 0.0)
 
-    return cond1 and cond2
+    return cond1 and cond2 and cond3 and cond4 and cond5 and cond6
 
 
 def _find_interval(t: jax.Array, boundaries: jax.Array):
@@ -488,6 +618,7 @@ def _draw_rxyz(
         config.randomized_interpolation_angle,
         config.range_of_motion_hinge if enable_range_of_motion else False,
         config.range_of_motion_hinge_method,
+        config.rom_halfsize,
         config.cdf_bins_min,
         config.cdf_bins_max,
         config.interpolation_method,
@@ -504,7 +635,11 @@ def _draw_pxyz(
     cor: bool = False,
 ) -> jax.Array:
     key_value, consume = jax.random.split(key_value)
-    POS_0 = jax.random.uniform(consume, minval=config.pos0_min, maxval=config.pos0_max)
+    POS_0 = jax.random.uniform(
+        consume,
+        minval=config.cor_pos0_min if cor else config.pos0_min,
+        maxval=config.cor_pos0_max if cor else config.pos0_max,
+    )
     max_iter = 100
     return _random.random_position_over_time(
         key_value,
@@ -590,10 +725,27 @@ def _draw_p3d_and_cor(
     __: jax.Array,
     cor: bool,
 ) -> jax.Array:
-    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, N, None, cor))(
-        jax.random.split(key_value, 3)
-    )
-    return pos.T
+    keys = jax.random.split(key_value, 3)
+
+    def draw(key, xyz: str):
+        return _draw_pxyz(
+            replace(
+                config,
+                pos_min=getattr(config, f"pos_min_p3d_{xyz}"),
+                pos_max=getattr(config, f"pos_max_p3d_{xyz}"),
+                pos0_min=getattr(config, f"pos0_min_p3d_{xyz}"),
+                pos0_max=getattr(config, f"pos0_max_p3d_{xyz}"),
+            ),
+            None,
+            key,
+            dt,
+            N,
+            None,
+            cor,
+        )[:, None]
+
+    px, py, pz = draw(keys[0], "x"), draw(keys[1], "y"), draw(keys[2], "z")
+    return jnp.concat((px, py, pz), axis=-1)
 
 
 def _draw_p3d(

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/ring/base.py

@@ -7,6 +7,7 @@ from jax.core import Tracer
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 import numpy as np
+import tree
 import tree_utils as tu
 
 import ring
@@ -590,6 +591,34 @@ class System(_Base):
 
         return sys
 
+    @staticmethod
+    def joint_type_simplification(typ: str) -> str:
+        if typ[:4] == "free":
+            if typ == "free_2d":
+                return "free_2d"
+            else:
+                return "free"
+        elif typ[:3] == "cor":
+            return "cor"
+        elif typ[:9] == "spherical":
+            return "spherical"
+        else:
+            return typ
+
+    @staticmethod
+    def joint_type_is_free_or_cor(typ: str) -> bool:
+        return System.joint_type_simplification(typ) in ["free", "cor"]
+
+    @staticmethod
+    def joint_type_is_spherical(typ: str) -> bool:
+        return System.joint_type_simplification(typ) == "spherical"
+
+    @staticmethod
+    def joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+        return System.joint_type_is_free_or_cor(typ) or System.joint_type_is_spherical(
+            typ
+        )
+
     def findall_imus(self, names: bool = True) -> list[str] | list[int]:
         bodies = [name for name in self.link_names if name[:3] == "imu"]
         return bodies if names else [self.name_to_idx(n) for n in bodies]
@@ -618,10 +647,20 @@ class System(_Base):
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
     def children(self, name: str, names: bool = False) -> list[int] | list[str]:
+        "List all direct children of body, does not include body itself"
         p = self.name_to_idx(name)
         bodies = [i for i in range(self.num_links()) if self.link_parents[i] == p]
         return bodies if (not names) else [self.idx_to_name(i) for i in bodies]
 
+    def findall_bodies_subsystem(
+        self, name: str, names: bool = False
+    ) -> list[int] | list[str]:
+        "List all children and children's children; does not include body itself"
+        children = self.children(name, names=True)
+        grandchildren = [self.findall_bodies_subsystem(n, names=True) for n in children]
+        bodies = tree.flatten([children, grandchildren])
+        return bodies if names else [self.name_to_idx(n) for n in bodies]
+
     def scan(self, f: Callable, in_types: str, *args, reverse: bool = False):
         """Scan `f` along each link in system whilst carrying along state.
 
@@ -889,7 +928,9 @@ def _parse_system(sys: System) -> System:
         assert d.size == a.size == s.size == qd_size, error_msg
         assert z.size == q_size, error_msg
 
-        if typ in ["spherical", "free", "cor"] and not isinstance(z, Tracer):
+        if System.joint_type_is_free_or_cor_or_spherical(typ) and not isinstance(
+            z, Tracer
+        ):
             assert jnp.allclose(
                 jnp.linalg.norm(z[:4]), 1.0
             ), f"not unit quat for link `{name}` of typ `{typ}` in model"
@@ -1030,7 +1071,7 @@ class State(_Base):
             def replace_by_unit_quat(_, idx_map, link_typ, link_idx):
                 nonlocal q
 
-                if link_typ in ["free", "cor", "spherical"]:
+                if sys.joint_type_is_free_or_cor_or_spherical(link_typ):
                     q_idxs_link = idx_map["q"](link_idx)
                     q = q.at[q_idxs_link.start].set(1.0)
 

{imt_ring-1.6.21 → imt_ring-1.6.23}/src/ring/io/xml/abstract.py

@@ -3,6 +3,7 @@ from typing import Tuple, TypeVar
 import jax
 import jax.numpy as jnp
 import numpy as np
+
 from ring import base
 
 T = TypeVar("T")
@@ -17,7 +18,7 @@ default_stiffness = lambda qd_size, **_: jnp.zeros((qd_size,))
 
 def default_zeropoint(q_size, link_typ: str, **_):
     zeropoint = jnp.zeros((q_size))
-    if link_typ in ["spherical", "free", "cor"]:
+    if base.System.joint_type_is_free_or_cor_or_spherical(link_typ):
         # zeropoint then is unit quaternion and not zeros
         zeropoint = zeropoint.at[0].set(1.0)
     return zeropoint