imt-ring 1.6.37__py3-none-any.whl → 1.6.39__py3-none-any.whl

ring/base.py

@@ -112,17 +112,71 @@ class _Base:
 
 @struct.dataclass
 class Transform(_Base):
-    """Represents the Transformation from Plücker A to Plücker B,
-    where B is located relative to A at `pos` in frame A and `rot` is the
-    relative quaternion from A to B.
-    Create using `Transform.create(pos=..., rot=...)
     """
+    Represents a spatial transformation between two coordinate frames using Plücker coordinates.
+
+    The `Transform` class defines the relative position and orientation of one frame (`B`)
+    with respect to another frame (`A`). The position (`pos`) is given in the coordinate frame
+    of `A`, and the rotation (`rot`) is expressed as a unit quaternion representing the relative
+    rotation from frame `A` to frame `B`.
+
+    Attributes:
+        pos (Vector):
+            The translation vector (position of `B` relative to `A`) in the coordinate frame of `A`.
+            Shape: `(..., 3)`, where `...` represents optional batch dimensions.
+        rot (Quaternion):
+            The unit quaternion representing the orientation of `B` relative to `A`.
+            Shape: `(..., 4)`, where `...` represents optional batch dimensions.
+
+    Methods:
+        create(pos: Optional[Vector] = None, rot: Optional[Quaternion] = None) -> Transform:
+            Creates a `Transform` instance with optional position and rotation.
+
+        zero(shape: Sequence[int] = ()) -> Transform:
+            Returns a zero transform with a given batch shape.
+
+        as_matrix() -> jax.Array:
+            Returns the 4x4 homogeneous transformation matrix representation of this transform.
+
+    Usage:
+        >>> pos = jnp.array([1.0, 2.0, 3.0])
+        >>> rot = jnp.array([1.0, 0.0, 0.0, 0.0])  # Identity quaternion
+        >>> T = Transform.create(pos, rot)
+        >>> print(T.pos)  # Output: [1. 2. 3.]
+        >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        >>> print(T.as_matrix())  # 4x4 transformation matrix
+    """  # noqa: E501
 
     pos: Vector
     rot: Quaternion
 
     @classmethod
     def create(cls, pos=None, rot=None):
+        """
+        Creates a `Transform` instance with the specified position and rotation.
+
+        At least one of `pos` or `rot` must be provided. If only `pos` is given, the rotation
+        defaults to the identity quaternion `[1, 0, 0, 0]`. If only `rot` is given, the position
+        defaults to `[0, 0, 0]`.
+
+        Args:
+            pos (Optional[Vector], default=None):
+                The position of frame `B` relative to frame `A`, expressed in frame `A` coordinates.
+                If `None`, defaults to a zero vector of shape `(3,)`.
+            rot (Optional[Quaternion], default=None):
+                The unit quaternion representing the orientation of `B` relative to `A`.
+                If `None`, defaults to the identity quaternion `(1, 0, 0, 0)`.
+
+        Returns:
+            Transform: A new `Transform` instance with the specified position and rotation.
+
+        Example:
+            >>> pos = jnp.array([1.0, 2.0, 3.0])
+            >>> rot = jnp.array([1.0, 0.0, 0.0, 0.0])  # Identity quaternion
+            >>> T = Transform.create(pos, rot)
+            >>> print(T.pos)  # Output: [1. 2. 3.]
+            >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        """  # noqa: E501
         assert not (pos is None and rot is None), "One must be given."
         shape_rot = rot.shape[:-1] if rot is not None else ()
         shape_pos = pos.shape[:-1] if pos is not None else ()
@@ -139,12 +193,49 @@ class Transform(_Base):
 
     @classmethod
     def zero(cls, shape=()) -> "Transform":
-        """Returns a zero transform with a batch shape."""
+        """
+        Returns a zero transform with a given batch shape.
+
+        This creates a transform with position `(0, 0, 0)` and an identity quaternion `(1, 0, 0, 0)`,
+        which represents no translation or rotation.
+
+        Args:
+            shape (Sequence[int], default=()):
+                The batch shape for the transform. Defaults to a scalar transform.
+
+        Returns:
+            Transform: A zero transform with the specified batch shape.
+
+        Example:
+            >>> T = Transform.zero()
+            >>> print(T.pos)  # Output: [0. 0. 0.]
+            >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        """  # noqa: E501
         pos = jnp.zeros(shape + (3,))
         rot = jnp.tile(jnp.array([1.0, 0.0, 0.0, 0.0]), shape + (1,))
         return Transform(pos, rot)
 
     def as_matrix(self) -> jax.Array:
+        """
+        Returns the 4x4 homogeneous transformation matrix representation of this transform.
+
+        The homogeneous transformation matrix is defined as:
+
+        ```
+        [ R  t ]
+        [ 0  1 ]
+        ```
+
+        where `R` is the 3x3 rotation matrix converted from the quaternion and `t` is the
+        3x1 position vector.
+
+        Returns:
+            jax.Array: A `(4, 4)` homogeneous transformation matrix.
+
+        Example:
+            >>> T = Transform.create(jnp.array([1.0, 2.0, 3.0]), jnp.array([1.0, 0.0, 0.0, 0.0]))
+            >>> print(T.as_matrix())  # Output: 4x4 matrix
+        """  # noqa: E501
         E = maths.quat_to_3x3(self.rot)
         return spatial.quadrants(aa=E, bb=E) @ spatial.xlt(self.pos)
 
@@ -402,7 +493,175 @@ QD_WIDTHS = {
 
 @struct.dataclass
 class System(_Base):
-    "System object. Create using `System.create(path_xml)`"
+    """
+    Represents a robotic system consisting of interconnected links and joints. Create it using `System.create(...)`
+
+    The `System` class models the kinematic and dynamic properties of a multibody
+    system, providing methods for state representation, transformations, joint
+    configuration management, and rendering. It supports both minimal and maximal
+    coordinate representations and can be parsed from or saved to XML files.
+
+    Attributes:
+        link_parents (list[int]):
+            A list specifying the parent index for each link. The root link has a parent index of `-1`.
+        links (Link):
+            A data structure containing information about all links in the system.
+        link_types (list[str]):
+            A list specifying the joint type for each link (e.g., "free", "hinge", "prismatic").
+        link_damping (jax.Array):
+            Joint damping coefficients for each link.
+        link_armature (jax.Array):
+            Armature inertia values for each joint.
+        link_spring_stiffness (jax.Array):
+            Stiffness values for joint springs.
+        link_spring_zeropoint (jax.Array):
+            Rest position of joint springs.
+        dt (float):
+            Simulation time step size.
+        geoms (list[Geometry]):
+            List of geometries associated with the system.
+        gravity (jax.Array):
+            Gravity vector applied to the system (default: `[0, 0, -9.81]`).
+        integration_method (str):
+            Integration method for simulation (default: "semi_implicit_euler").
+        mass_mat_iters (int):
+            Number of iterations for mass matrix calculations.
+        link_names (list[str]):
+            Names of the links in the system.
+        model_name (Optional[str]):
+            Name of the system model (if available).
+        omc (list[MaxCoordOMC | None]):
+            List of optional Maximal Coordinate representations.
+
+    Methods:
+        num_links() -> int:
+            Returns the number of links in the system.
+
+        q_size() -> int:
+            Returns the total number of generalized coordinates (`q`) in the system.
+
+        qd_size() -> int:
+            Returns the total number of generalized velocities (`qd`) in the system.
+
+        name_to_idx(name: str) -> int:
+            Returns the index of a link given its name.
+
+        idx_to_name(idx: int, allow_world: bool = False) -> str:
+            Returns the name of a link given its index. If `allow_world` is `True`,
+            returns `"world"` for index `-1`.
+
+        idx_map(type: str) -> dict:
+            Returns a dictionary mapping link names to their indices for a specified type
+            (`"l"`, `"q"`, or `"d"`).
+
+        parent_name(name: str) -> str:
+            Returns the name of the parent link for a given link.
+
+        change_model_name(new_name: Optional[str] = None, prefix: Optional[str] = None, suffix: Optional[str] = None) -> "System":
+            Changes the name of the system model.
+
+        change_link_name(old_name: str, new_name: str) -> "System":
+            Renames a specific link.
+
+        add_prefix_suffix(prefix: Optional[str] = None, suffix: Optional[str] = None) -> "System":
+            Adds both a prefix and suffix to all link names.
+
+        freeze(name: str | list[str]) -> "System":
+            Freezes the specified link(s), making them immovable.
+
+        unfreeze(name: str, new_joint_type: str) -> "System":
+            Unfreezes a frozen link and assigns it a new joint type.
+
+        change_joint_type(name: str, new_joint_type: str, **kwargs) -> "System":
+            Changes the joint type of a specified link.
+
+        joint_type_simplification(typ: str) -> str:
+            Returns a simplified representation of the given joint type.
+
+        joint_type_is_free_or_cor(typ: str) -> bool:
+            Checks if a joint type is either "free" or "cor".
+
+        joint_type_is_spherical(typ: str) -> bool:
+            Checks if a joint type is "spherical".
+
+        joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+            Checks if a joint type is "free", "cor", or "spherical".
+
+        findall_imus(names: bool = True) -> list[str] | list[int]:
+            Finds all IMU sensors in the system.
+
+        findall_segments(names: bool = True) -> list[str] | list[int]:
+            Finds all non-IMU segments in the system.
+
+        findall_bodies_to_world(names: bool = False) -> list[int] | list[str]:
+            Returns all bodies directly connected to the world.
+
+        find_body_to_world(name: bool = False) -> int | str:
+            Returns the root body connected to the world.
+
+        findall_bodies_with_jointtype(typ: str, names: bool = False) -> list[int] | list[str]:
+            Returns all bodies with the specified joint type.
+
+        children(name: str, names: bool = False) -> list[int] | list[str]:
+            Returns the direct children of a given body.
+
+        findall_bodies_subsystem(name: str, names: bool = False) -> list[int] | list[str]:
+            Finds all bodies in the subsystem rooted at a given link.
+
+        scan(f: Callable, in_types: str, *args, reverse: bool = False):
+            Iterates over system elements while applying a function.
+
+        parse() -> "System":
+            Parses the system, performing consistency checks and computing spatial inertia tensors.
+
+        render(xs: Optional[Transform | list[Transform]] = None, **kwargs) -> list[np.ndarray]:
+            Renders frames of the system using maximal coordinates.
+
+        render_prediction(xs: Transform | list[Transform], yhat: dict | jax.Array | np.ndarray, **kwargs):
+            Renders a predicted state transformation.
+
+        delete_system(link_name: str | list[str], strict: bool = True):
+            Removes a subsystem from the system.
+
+        make_sys_noimu(imu_link_names: Optional[list[str]] = None):
+            Returns a version of the system without IMU sensors.
+
+        inject_system(other_system: "System", at_body: Optional[str] = None):
+            Merges another system into this one.
+
+        morph_system(new_parents: Optional[list[int | str]] = None, new_anchor: Optional[int | str] = None):
+            Reorders the system’s link hierarchy.
+
+        from_xml(path: str, seed: int = 1) -> "System":
+            Loads a system from an XML file.
+
+        from_str(xml: str, seed: int = 1) -> "System":
+            Loads a system from an XML string.
+
+        to_str(warn: bool = True) -> str:
+            Serializes the system to an XML string.
+
+        to_xml(path: str) -> None:
+            Saves the system as an XML file.
+
+        create(path_or_str: str, seed: int = 1) -> "System":
+            Creates a `System` instance from an XML file or string.
+
+        coordinate_vector_to_q(q: jax.Array, custom_joints: dict[str, Callable] = {}) -> jax.Array:
+            Converts a coordinate vector to minimal coordinates (`q`), applying
+            constraints such as quaternion normalization.
+
+    Raises:
+        AssertionError: If the system structure is invalid (e.g., duplicate link names, incorrect parent-child relationships).
+        InvalidSystemError: If an operation results in an inconsistent system state.
+
+    Notes:
+        - The system must be parsed before use to ensure consistency.
+        - The system supports batch operations using JAX for efficient computations.
+        - Joint types include revolute ("rx", "ry", "rz"), prismatic ("px", "py", "pz"), spherical, free, and more.
+        - Inertial properties of links are computed automatically from associated geometries.
+    """  # noqa: E501
+
     link_parents: list[int] = struct.field(False)
     links: Link
     link_types: list[str] = struct.field(False)
@@ -429,25 +688,32 @@ class System(_Base):
     omc: list[MaxCoordOMC | None] = struct.field(True, default_factory=lambda: [])
 
     def num_links(self) -> int:
+        "Returns the number of links in the system."
         return len(self.link_parents)
 
     def q_size(self) -> int:
+        "Returns the total number of generalized coordinates (`q`) in the system."
         return sum([Q_WIDTHS[typ] for typ in self.link_types])
 
     def qd_size(self) -> int:
+        "Returns the total number of generalized velocities (`qd`) in the system."
         return sum([QD_WIDTHS[typ] for typ in self.link_types])
 
     def name_to_idx(self, name: str) -> int:
+        "Returns the index of a link given its name."
         return self.link_names.index(name)
 
     def idx_to_name(self, idx: int, allow_world: bool = False) -> str:
+        """Returns the name of a link given its index. If `allow_world` is `True`,
+        returns `"world"` for index `-1`."""
         if allow_world and idx == -1:
             return "world"
         assert idx >= 0, "Worldbody index has no name."
         return self.link_names[idx]
 
     def idx_map(self, type: str) -> dict:
-        "type: is either `l` or `q` or `d`"
+        """Returns a dictionary mapping link names to their indices for a specified type
+        (`"l"`, `"q"`, or `"d"`)."""
         dict_int_slices = {}
 
         def f(_, idx_map, name: str, link_idx: int):
@@ -458,10 +724,11 @@ class System(_Base):
         return dict_int_slices
 
     def parent_name(self, name: str) -> str:
+        "Returns the name of the parent link for a given link."
         return self.idx_to_name(self.link_parents[self.name_to_idx(name)])
 
     def add_prefix(self, prefix: str = "") -> "System":
-        return self.replace(link_names=[prefix + name for name in self.link_names])
+        return self.add_prefix_suffix(prefix=prefix)
 
     def change_model_name(
         self,
@@ -469,6 +736,7 @@ class System(_Base):
         prefix: Optional[str] = None,
         suffix: Optional[str] = None,
     ) -> "System":
+        "Changes the name of the system model."
         if prefix is None:
             prefix = ""
         if suffix is None:
@@ -479,6 +747,7 @@ class System(_Base):
         return self.replace(model_name=name)
 
     def change_link_name(self, old_name: str, new_name: str) -> "System":
+        "Renames a specific link."
         old_idx = self.name_to_idx(old_name)
         new_link_names = self.link_names.copy()
         new_link_names[old_idx] = new_name
@@ -487,6 +756,7 @@ class System(_Base):
     def add_prefix_suffix(
         self, prefix: Optional[str] = None, suffix: Optional[str] = None
     ) -> "System":
+        "Adds either or, or both a prefix and suffix to all link names."
         if prefix is None:
             prefix = ""
         if suffix is None:
@@ -526,6 +796,7 @@ class System(_Base):
         return _update_sys_if_replace_joint_type(self, logic_replace_free_with_cor)
 
     def freeze(self, name: str | list[str]):
+        "Freezes the specified link(s), making them immovable (uses `frozen` joint)"
         if isinstance(name, list):
             sys = self
             for n in name:
@@ -544,6 +815,7 @@ class System(_Base):
         return _update_sys_if_replace_joint_type(self, logic_freeze)
 
     def unfreeze(self, name: str, new_joint_type: str):
+        "Unfreezes a frozen link and assigns it a new joint type."
         assert self.link_types[self.name_to_idx(name)] == "frozen"
         assert new_joint_type != "frozen"
 
@@ -560,7 +832,8 @@ class System(_Base):
         seed: int = 1,
         warn: bool = True,
     ):
-        "By default damping, stiffness are set to zero."
+        """Changes the joint type of a specified link.
+        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]
@@ -594,6 +867,7 @@ class System(_Base):
 
     @staticmethod
     def joint_type_simplification(typ: str) -> str:
+        "Returns a simplified name of the given joint type."
         if typ[:4] == "free":
             if typ == "free_2d":
                 return "free_2d"
@@ -608,23 +882,28 @@ class System(_Base):
 
     @staticmethod
     def joint_type_is_free_or_cor(typ: str) -> bool:
+        'Checks if a joint type is either "free" or "cor".'
         return System.joint_type_simplification(typ) in ["free", "cor"]
 
     @staticmethod
     def joint_type_is_spherical(typ: str) -> bool:
+        'Checks if a joint type is "spherical".'
         return System.joint_type_simplification(typ) == "spherical"
 
     @staticmethod
     def joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+        'Checks if a joint type is "free", "cor", or "spherical".'
         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]:
+        "Finds all IMU sensors in the system."
         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]
 
     def findall_segments(self, names: bool = True) -> list[str] | list[int]:
+        "Finds all non-IMU segments in the system."
         imus = self.findall_imus(names=True)
         bodies = [name for name in self.link_names if name not in imus]
         return bodies if names else [self.name_to_idx(n) for n in bodies]
@@ -633,10 +912,12 @@ class System(_Base):
         return [self.idx_to_name(i) for i in bodies]
 
     def findall_bodies_to_world(self, names: bool = False) -> list[int] | list[str]:
+        "Returns all bodies directly connected to the world."
         bodies = [i for i, p in enumerate(self.link_parents) if p == -1]
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
     def find_body_to_world(self, name: bool = False) -> int | str:
+        "Returns the root body connected to the world."
         bodies = self.findall_bodies_to_world(names=name)
         assert len(bodies) == 1
         return bodies[0]
@@ -644,6 +925,7 @@ class System(_Base):
     def findall_bodies_with_jointtype(
         self, typ: str, names: bool = False
     ) -> list[int] | list[str]:
+        "Returns all bodies with the specified joint type."
         bodies = [i for i, _typ in enumerate(self.link_types) if _typ == typ]
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
@@ -781,20 +1063,25 @@ class System(_Base):
 
     @staticmethod
     def from_xml(path: str, seed: int = 1):
+        "Loads a system from an XML file."
         return ring.io.load_sys_from_xml(path, seed)
 
     @staticmethod
     def from_str(xml: str, seed: int = 1):
+        "Loads a system from an XML string."
         return ring.io.load_sys_from_str(xml, seed)
 
     def to_str(self, warn: bool = True) -> str:
+        "Serializes the system to an XML string."
         return ring.io.save_sys_to_str(self, warn=warn)
 
     def to_xml(self, path: str) -> None:
+        "Saves the system to an XML file."
         ring.io.save_sys_to_xml(self, path)
 
     @classmethod
     def create(cls, path_or_str: str, seed: int = 1) -> "System":
+        "Creates a `System` instance from an XML file or string."
         path = Path(path_or_str).with_suffix(".xml")
 
         exists = False
@@ -807,14 +1094,15 @@ class System(_Base):
         if exists:
             return cls.from_xml(path, seed=seed)
         else:
-            return cls.from_str(path_or_str)
+            return cls.from_str(path_or_str, seed=seed)
 
     def coordinate_vector_to_q(
         self,
         q: jax.Array,
         custom_joints: dict[str, Callable] = {},
     ) -> jax.Array:
-        """Map a coordinate vector `q` to the minimal coordinates vector of the sys"""
+        """Converts a coordinate vector to minimal coordinates (`q`), applying
+        constraints such as quaternion normalization."""
         # Does, e.g.
         # - normalize quaternions
         # - hinge joints in [-pi, pi]
@@ -1026,14 +1314,37 @@ def _scan_sys(sys: System, f: Callable, in_types: str, *args, reverse: bool = Fa
 
 @struct.dataclass
 class State(_Base):
-    """The static and dynamic state of a system in minimal and maximal coordinates.
-    Use `.create()` to create this object.
-
-    Args:
-        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.
     """
+    Represents the state of a dynamic system in minimal and maximal coordinates.
+
+    The `State` class encapsulates both the configuration (`q`) and velocity (`qd`)
+    of the system in minimal coordinates, as well as the corresponding transforms (`x`)
+    in maximal coordinates.
+
+    Attributes:
+        q (jax.Array):
+            The joint positions (generalized coordinates) of the system. The size
+            of `q` matches `sys.q_size()`.
+        qd (jax.Array):
+            The joint velocities (generalized velocities) of the system. The size
+            of `qd` matches `sys.qd_size()`.
+        x (Transform):
+            The maximal coordinate representation of all system links, expressed as
+            a `Transform` object.
+
+    Methods:
+        create(sys: System, q: Optional[jax.Array] = None,
+               qd: Optional[jax.Array] = None, x: Optional[Transform] = None,
+               key: Optional[jax.Array] = None,
+               custom_joints: dict[str, Callable] = {}) -> State:
+            Creates a `State` instance for a given system with optional initial conditions.
+
+    Usage:
+        >>> sys = System.create("model.xml")
+        >>> state = State.create(sys)
+        >>> print(state.q.shape)  # Should match sys.q_size()
+        >>> print(state.qd.shape)  # Should match sys.qd_size()
+    """  # noqa: E501
 
     q: jax.Array
     qd: jax.Array
@@ -1048,19 +1359,37 @@ class State(_Base):
         x: Optional[Transform] = None,
         key: Optional[jax.Array] = None,
         custom_joints: dict[str, Callable] = {},
-    ):
-        """Create state of system.
+    ) -> "State":
+        """
+        Creates a `State` instance for the given system with optional initial conditions.
+
+        If no initial values are provided, joint positions (`q`) and velocities (`qd`)
+        are initialized to zero, except for free and spherical joints, which have unit quaternions.
 
         Args:
-            sys (System): The system for which to create a state.
-            q (jax.Array, optional): The joint values of the system. Defaults to None.
-            Which then defaults to zeros.
-            qd (jax.Array, optional): The joint velocities of the system.
-            Defaults to None. Which then defaults to zeros.
+            sys (System):
+                The system for which to create a state.
+            q (Optional[jax.Array], default=None):
+                Initial joint positions. If `None`, defaults to zeros, with unit quaternion initialization
+                for free and spherical joints.
+            qd (Optional[jax.Array], default=None):
+                Initial joint velocities. If `None`, defaults to zeros.
+            x (Optional[Transform], default=None):
+                Initial maximal coordinates of the system links. If `None`, defaults to zero transforms.
+            key (Optional[jax.Array], default=None):
+                Random key for initializing `q` if no values are provided.
+            custom_joints (dict[str, Callable], default={}):
+                Custom joint functions for mapping coordinate vectors to minimal coordinates.
 
         Returns:
-            (State): Create State object.
-        """
+            State: A new instance of the `State` class representing the initialized system state.
+
+        Example:
+            >>> sys = System.create("model.xml")
+            >>> state = State.create(sys)
+            >>> print(state.q.shape)  # Should match sys.q_size()
+            >>> print(state.qd.shape)  # Should match sys.qd_size()
+        """  # noqa: E501
         if key is not None:
             assert q is None
             q = jax.random.normal(key, shape=(sys.q_size(),))

ring/io/xml/from_xml.py

@@ -252,7 +252,7 @@ def load_sys_from_str(xml_str: str, seed: int = 1) -> base.System:
 
     # numpy -> jax
     # we load using numpy in order to have float64 precision
-    sys = jax.tree_map(jax.numpy.asarray, sys)
+    sys = jax.tree.map(jax.numpy.asarray, sys)
 
     sys = jcalc._init_joint_params(jax.random.PRNGKey(seed), sys)
 

ring/ml/base.py

@@ -13,13 +13,13 @@ from ring.utils import pickle_save
 def _to_3d(tree):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[None], tree)
+    return jax.tree.map(lambda arr: arr[None], tree)
 
 
 def _to_2d(tree, i: int = 0):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[i], tree)
+    return jax.tree.map(lambda arr: arr[i], tree)
 
 
 class AbstractFilter(ABC):
@@ -297,7 +297,8 @@ class NoGraph_FilterWrapper(AbstractFilterWrapper):
 
         if self._quat_normalize:
             assert yhat.shape[-1] == 4, f"yhat.shape={yhat.shape}"
-            yhat = ring.maths.safe_normalize(yhat)
+            # yhat = ring.maths.safe_normalize(yhat)
+            yhat = yhat / jnp.linalg.norm(yhat, axis=-1, keepdims=True)
 
         return yhat, state
 

ring/ml/ml_utils.py

@@ -161,7 +161,7 @@ def _flatten_convert_filter_nested_dict(
     metrices: NestedDict, filter_nan_inf: bool = True
 ):
     metrices = _flatten_dict(metrices)
-    metrices = jax.tree_map(_to_float_if_not_string, metrices)
+    metrices = jax.tree.map(_to_float_if_not_string, metrices)
 
     if not filter_nan_inf:
         return metrices
@@ -216,7 +216,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     from jax.experimental import jax2tf
     import tensorflow as tf
 
-    signature = jax.tree_map(
+    signature = jax.tree.map(
         lambda arr: tf.TensorSpec(list(arr.shape), tf.float32), input
     )
 
@@ -241,7 +241,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     if validate:
         output_jax = jax_func(*input)
         output_tf = tf.saved_model.load(path)(*input)
-        jax.tree_map(
+        jax.tree.map(
             lambda a1, a2: np.allclose(a1, a2, atol=1e-5, rtol=1e-5),
             output_jax,
             output_tf,

ring/ml/ringnet.py

@@ -248,7 +248,7 @@ class RING(ml_base.AbstractFilter):
         params, state = self.forward_lam_factory(lam=lam).init(key, X)
 
         if bs is not None:
-            state = jax.tree_map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
+            state = jax.tree.map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
 
         return params, state
 

ring/ml/train.py

@@ -50,7 +50,7 @@ def _build_step_fn(
         # this vmap maps along batch-axis, not time-axis
         # time-axis is handled by `metric_fn`
         pipe = lambda q, qhat: jnp.mean(jax.vmap(metric_fn)(q, qhat))
-        error_tree = jax.tree_map(pipe, y, yhat)
+        error_tree = jax.tree.map(pipe, y, yhat)
         return jnp.mean(tree_utils.batch_concat(error_tree, 0)), state
 
     @partial(
@@ -274,7 +274,7 @@ def _build_eval_fn(
             ), f"The metric identitifier {metric_name} is not unique"
 
             pipe = lambda q, qhat: reduce_fn(jax.vmap(jax.vmap(metric_fn))(q, qhat))
-            values.update({metric_name: jax.tree_map(pipe, y, yhat)})
+            values.update({metric_name: jax.tree.map(pipe, y, yhat)})
 
         return values