jaxsim 0.5.1.dev126__py3-none-any.whl → 0.5.1.dev139__py3-none-any.whl

jaxsim/integrators/variable_step.py

@@ -1,3 +1,4 @@
+import dataclasses
 import functools
 from typing import Any, ClassVar, Generic
 
@@ -216,6 +217,17 @@ def local_error_estimation(
 
 @jax_dataclasses.pytree_dataclass
 class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    An Embedded Runge-Kutta integrator.
+
+    This class implements a general-purpose Embedded Runge-Kutta integrator
+    that can be used to solve ordinary differential equations with adaptive
+    step sizes.
+
+    The integrator is based on an Explicit Runge-Kutta method, and it uses
+    two different solutions to estimate the local integration error. The
+    error is then used to adapt the step size to reach a desired accuracy.
+    """
 
     AfterInitKey: ClassVar[str] = "after_init"
     InitializingKey: ClassVar[str] = "initializing"
@@ -243,6 +255,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
     # Maximum number of rejected steps when the Δt needs to be reduced.
     max_step_rejections: Static[jtp.IntLike] = MAX_STEP_REJECTIONS_DEFAULT
 
+    index_of_fsal: jtp.IntLike | None = None
+    fsal_enabled_if_supported: bool = False
+
     def init(
         self,
         x0: State,
@@ -257,6 +272,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             x0: The initial state of the system.
             t0: The initial time of the system.
             dt: The time step of the integration.
+            **kwargs: Additional parameters.
 
         Returns:
             The metadata of the integrator to be passed to the first step.
@@ -296,6 +312,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Integrate the system for a single step.
+        """
 
         # This method is called differently in three stages:
         #
@@ -512,10 +531,16 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
     @property
     def order_of_solution(self) -> int:
+        """
+        The order of the solution.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution]
 
     @property
     def order_of_solution_estimate(self) -> int:
+        """
+        The order of the solution estimate.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution_estimate]
 
     @classmethod
@@ -534,17 +559,36 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         max_step_rejections: jtp.IntLike = MAX_STEP_REJECTIONS_DEFAULT,
         **kwargs,
     ) -> Self:
+        """
+        Build an Embedded Runge-Kutta integrator.
+
+        Args:
+            dynamics: The system dynamics function.
+            fsal_enabled_if_supported:
+                Whether to enable the FSAL property if supported by the integrator.
+            dt_max: The maximum step size.
+            dt_min: The minimum step size.
+            rtol: The relative tolerance.
+            atol: The absolute tolerance.
+            safety: The safety factor to shrink the step size.
+            beta_max: The maximum factor to increase the step size.
+            beta_min: The minimum factor to increase the step size.
+            max_step_rejections: The maximum number of step rejections.
+            **kwargs: Additional parameters.
+        """
+
+        b = cls.__dataclass_fields__["b"].default_factory()
 
         # Check that b.T has enough rows based on the configured index of the
         # solution estimate. This is necessary for embedded methods.
         if (
             cls.row_index_of_solution_estimate is not None
-            and cls.row_index_of_solution_estimate >= cls.b.T.shape[0]
+            and cls.row_index_of_solution_estimate >= b.T.shape[0]
         ):
             msg = "The index of the solution estimate ({}-th row of `b.T`) "
             msg += "is out of range ({})."
             raise ValueError(
-                msg.format(cls.row_index_of_solution_estimate, cls.b.T.shape[0])
+                msg.format(cls.row_index_of_solution_estimate, b.T.shape[0])
             )
 
         integrator = super().build(
@@ -569,65 +613,92 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class HeunEulerSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mixin):
+    """
+    The Heun-Euler integrator for SO(3) dynamics.
+    """
 
-    A: ClassVar[jtp.Matrix] = jnp.array(
-        [
-            [0, 0],
-            [1, 0],
-        ]
-    ).astype(float)
-
-    b: ClassVar[jtp.Matrix] = (
-        jnp.atleast_2d(
-            jnp.array(
-                [
-                    [1 / 2, 1 / 2],
-                    [1, 0],
-                ]
-            ),
-        )
-        .astype(float)
-        .transpose()
+    A: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [
+                [0, 0],
+                [1, 0],
+            ]
+        ).astype(float),
+        compare=False,
+    )
+
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: (
+            jnp.atleast_2d(
+                jnp.array(
+                    [
+                        [1 / 2, 1 / 2],
+                        [1, 0],
+                    ]
+                ),
+            )
+            .astype(float)
+            .transpose()
+        ),
+        compare=False,
     )
 
-    c: ClassVar[jtp.Vector] = jnp.array(
-        [0, 1],
-    ).astype(float)
+    c: jtp.Vector = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [0, 1],
+        ).astype(float),
+        compare=False,
+    )
 
     row_index_of_solution: ClassVar[int] = 0
     row_index_of_solution_estimate: ClassVar[int | None] = 1
 
     order_of_bT_rows: ClassVar[tuple[int, ...]] = (2, 1)
 
+    index_of_fsal: jtp.IntLike | None = None
+    fsal_enabled_if_supported: bool = False
+
 
 @jax_dataclasses.pytree_dataclass
 class BogackiShampineSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mixin):
+    """
+    The Bogacki-Shampine integrator for SO(3) dynamics.
+    """
 
-    A: ClassVar[jtp.Matrix] = jnp.array(
-        [
-            [0, 0, 0, 0],
-            [1 / 2, 0, 0, 0],
-            [0, 3 / 4, 0, 0],
-            [2 / 9, 1 / 3, 4 / 9, 0],
-        ]
-    ).astype(float)
-
-    b: ClassVar[jtp.Matrix] = (
-        jnp.atleast_2d(
-            jnp.array(
-                [
-                    [2 / 9, 1 / 3, 4 / 9, 0],
-                    [7 / 24, 1 / 4, 1 / 3, 1 / 8],
-                ]
-            ),
-        )
-        .astype(float)
-        .transpose()
+    A: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [
+                [0, 0, 0, 0],
+                [1 / 2, 0, 0, 0],
+                [0, 3 / 4, 0, 0],
+                [2 / 9, 1 / 3, 4 / 9, 0],
+            ]
+        ).astype(float),
+        compare=False,
     )
 
-    c: ClassVar[jtp.Vector] = jnp.array(
-        [0, 1 / 2, 3 / 4, 1],
-    ).astype(float)
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: (
+            jnp.atleast_2d(
+                jnp.array(
+                    [
+                        [2 / 9, 1 / 3, 4 / 9, 0],
+                        [7 / 24, 1 / 4, 1 / 3, 1 / 8],
+                    ]
+                ),
+            )
+            .astype(float)
+            .transpose()
+        ),
+        compare=False,
+    )
+
+    c: jtp.Vector = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [0, 1 / 2, 3 / 4, 1],
+        ).astype(float),
+        compare=False,
+    )
 
     row_index_of_solution: ClassVar[int] = 0
     row_index_of_solution_estimate: ClassVar[int | None] = 1

jaxsim/math/adjoint.py

@@ -7,10 +7,14 @@ from .skew import Skew
 
 
 class Adjoint:
+    """
+    A utility class for adjoint matrix operations.
+    """
+
     @staticmethod
     def from_quaternion_and_translation(
-        quaternion: jtp.Vector = jnp.array([1.0, 0, 0, 0]),
-        translation: jtp.Vector = jnp.zeros(3),
+        quaternion: jtp.Vector | None = None,
+        translation: jtp.Vector | None = None,
         inverse: bool = False,
         normalize_quaternion: bool = False,
     ) -> jtp.Matrix:
@@ -18,8 +22,8 @@ class Adjoint:
         Create an adjoint matrix from a quaternion and a translation.
 
         Args:
-            quaternion (jtp.Vector): A quaternion vector (4D) representing orientation.
-            translation (jtp.Vector): A translation vector (3D).
+            quaternion (jtp.Vector): A quaternion vector (4D) representing orientation. Default is [1, 0, 0, 0].
+            translation (jtp.Vector): A translation vector (3D). Default is [0, 0, 0].
             inverse (bool): Whether to compute the inverse adjoint. Default is False.
             normalize_quaternion (bool): Whether to normalize the quaternion before creating the adjoint.
                                          Default is False.
@@ -27,6 +31,8 @@ class Adjoint:
         Returns:
             jtp.Matrix: The adjoint matrix.
         """
+        quaternion = quaternion if quaternion is not None else jnp.array([1.0, 0, 0, 0])
+        translation = translation if translation is not None else jnp.zeros(3)
         assert quaternion.size == 4
         assert translation.size == 3
 
@@ -61,21 +67,24 @@ class Adjoint:
 
     @staticmethod
     def from_rotation_and_translation(
-        rotation: jtp.Matrix = jnp.eye(3),
-        translation: jtp.Vector = jnp.zeros(3),
+        rotation: jtp.Matrix | None = None,
+        translation: jtp.Vector | None = None,
         inverse: bool = False,
     ) -> jtp.Matrix:
         """
         Create an adjoint matrix from a rotation matrix and a translation vector.
 
         Args:
-            rotation (jtp.Matrix): A 3x3 rotation matrix.
-            translation (jtp.Vector): A translation vector (3D).
+            rotation (jtp.Matrix): A 3x3 rotation matrix. Default is identity.
+            translation (jtp.Vector): A translation vector (3D). Default is [0, 0, 0].
             inverse (bool): Whether to compute the inverse adjoint. Default is False.
 
         Returns:
             jtp.Matrix: The adjoint matrix.
         """
+        rotation = rotation if rotation is not None else jnp.eye(3)
+        translation = translation if translation is not None else jnp.zeros(3)
+
         assert rotation.shape == (3, 3)
         assert translation.size == 3
 
@@ -105,7 +114,7 @@ class Adjoint:
         Convert an adjoint matrix to a transformation matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix (6x6).
+            adjoint: The adjoint matrix (6x6).
 
         Returns:
             jtp.Matrix: The transformation matrix (4x4).
@@ -131,7 +140,7 @@ class Adjoint:
         Compute the inverse of an adjoint matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix.
+            adjoint: The adjoint matrix.
 
         Returns:
             jtp.Matrix: The inverse adjoint matrix.

jaxsim/math/cross.py

@@ -6,13 +6,17 @@ from .skew import Skew
 
 
 class Cross:
+    """
+    A utility class for cross product matrix operations.
+    """
+
     @staticmethod
     def vx(velocity_sixd: jtp.Vector) -> jtp.Matrix:
         """
         Compute the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
             jtp.Matrix: The cross product matrix (6x6).
@@ -37,7 +41,7 @@ class Cross:
         Compute the negative transpose of the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
             jtp.Matrix: The negative transpose of the cross product matrix (6x6).

jaxsim/math/inertia.py

@@ -6,15 +6,19 @@ from .skew import Skew
 
 
 class Inertia:
+    """
+    A utility class for inertia matrix operations.
+    """
+
     @staticmethod
     def to_sixd(mass: jtp.Float, com: jtp.Vector, I: jtp.Matrix) -> jtp.Matrix:
         """
         Convert mass, center of mass, and inertia matrix to a 6x6 inertia matrix.
 
         Args:
-            mass (jtp.Float): The mass of the body.
-            com (jtp.Vector): The center of mass position (3D).
-            I (jtp.Matrix): The 3x3 inertia matrix.
+            mass: The mass of the body.
+            com: The center of mass position (3D).
+            I: The 3x3 inertia matrix.
 
         Returns:
             jtp.Matrix: The 6x6 inertia matrix.
@@ -42,7 +46,7 @@ class Inertia:
         Convert a 6x6 inertia matrix to mass, center of mass, and inertia matrix.
 
         Args:
-            M (jtp.Matrix): The 6x6 inertia matrix.
+            M: The 6x6 inertia matrix.
 
         Returns:
             tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).

jaxsim/math/quaternion.py

@@ -8,13 +8,17 @@ from .utils import safe_norm
 
 
 class Quaternion:
+    """
+    A utility class for quaternion operations.
+    """
+
     @staticmethod
     def to_xyzw(wxyz: jtp.Vector) -> jtp.Vector:
         """
         Convert a quaternion from WXYZ to XYZW representation.
 
         Args:
-            wxyz (jtp.Vector): Quaternion in WXYZ representation.
+            wxyz: Quaternion in WXYZ representation.
 
         Returns:
             jtp.Vector: Quaternion in XYZW representation.
@@ -27,7 +31,7 @@ class Quaternion:
         Convert a quaternion from XYZW to WXYZ representation.
 
         Args:
-            xyzw (jtp.Vector): Quaternion in XYZW representation.
+            xyzw: Quaternion in XYZW representation.
 
         Returns:
             jtp.Vector: Quaternion in WXYZ representation.
@@ -40,7 +44,7 @@ class Quaternion:
         Convert a quaternion to a direction cosine matrix (DCM).
 
         Args:
-            quaternion (jtp.Vector): Quaternion in XYZW representation.
+            quaternion: Quaternion in XYZW representation.
 
         Returns:
             jtp.Matrix: Direction cosine matrix (DCM).
@@ -53,7 +57,7 @@ class Quaternion:
         Convert a direction cosine matrix (DCM) to a quaternion.
 
         Args:
-            dcm (jtp.Matrix): Direction cosine matrix (DCM).
+            dcm: Direction cosine matrix (DCM).
 
         Returns:
             jtp.Vector: Quaternion in XYZW representation.
@@ -71,8 +75,8 @@ class Quaternion:
         Compute the derivative of a quaternion given angular velocity.
 
         Args:
-            quaternion (jtp.Vector): Quaternion in XYZW representation.
-            omega (jtp.Vector): Angular velocity vector.
+            quaternion: Quaternion in XYZW representation.
+            omega: Angular velocity vector.
             omega_in_body_fixed (bool): Whether the angular velocity is in the body-fixed frame.
             K (float): A scaling factor.
 

jaxsim/math/rotation.py

@@ -8,6 +8,9 @@ from .utils import safe_norm
 
 
 class Rotation:
+    """
+    A utility class for rotation matrix operations.
+    """
 
     @staticmethod
     def x(theta: jtp.Float) -> jtp.Matrix:
@@ -15,7 +18,7 @@ class Rotation:
         Generate a 3D rotation matrix around the X-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.
@@ -29,7 +32,7 @@ class Rotation:
         Generate a 3D rotation matrix around the Y-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.
@@ -43,7 +46,7 @@ class Rotation:
         Generate a 3D rotation matrix around the Z-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.

jaxsim/math/skew.py

@@ -14,7 +14,7 @@ class Skew:
         Compute the skew-symmetric matrix (wedge operator) of a 3D vector.
 
         Args:
-            vector (jtp.Vector): A 3D vector.
+            vector: A 3D vector.
 
         Returns:
             jtp.Matrix: The skew-symmetric matrix corresponding to the input vector.
@@ -31,7 +31,7 @@ class Skew:
         Extract the 3D vector from a skew-symmetric matrix (vee operator).
 
         Args:
-            matrix (jtp.Matrix): A 3x3 skew-symmetric matrix.
+            matrix: A 3x3 skew-symmetric matrix.
 
         Returns:
             jtp.Vector: The 3D vector extracted from the input matrix.

jaxsim/math/transform.py

@@ -5,11 +5,14 @@ import jaxsim.typing as jtp
 
 
 class Transform:
+    """
+    A utility class for transformation matrix operations.
+    """
 
     @staticmethod
     def from_quaternion_and_translation(
-        quaternion: jtp.VectorLike = jnp.array([1.0, 0, 0, 0]),
-        translation: jtp.VectorLike = jnp.zeros(3),
+        quaternion: jtp.VectorLike | None = None,
+        translation: jtp.VectorLike | None = None,
         inverse: jtp.BoolLike = False,
         normalize_quaternion: jtp.BoolLike = False,
     ) -> jtp.Matrix:
@@ -27,6 +30,9 @@ class Transform:
             The 4x4 transformation matrix representing the SE(3) transformation.
         """
 
+        quaternion = quaternion if quaternion is not None else jnp.array([1.0, 0, 0, 0])
+        translation = translation if translation is not None else jnp.zeros(3)
+
         W_Q_B = jnp.array(quaternion).astype(float)
         W_p_B = jnp.array(translation).astype(float)
 
@@ -44,8 +50,8 @@ class Transform:
 
     @staticmethod
     def from_rotation_and_translation(
-        rotation: jtp.MatrixLike = jnp.eye(3),
-        translation: jtp.VectorLike = jnp.zeros(3),
+        rotation: jtp.MatrixLike | None = None,
+        translation: jtp.VectorLike | None = None,
         inverse: jtp.BoolLike = False,
     ) -> jtp.Matrix:
         """
@@ -59,6 +65,8 @@ class Transform:
         Returns:
             The 4x4 transformation matrix representing the SE(3) transformation.
         """
+        rotation = rotation if rotation is not None else jnp.eye(3)
+        translation = translation if translation is not None else jnp.zeros(3)
 
         A_R_B = jnp.array(rotation).astype(float)
         W_p_B = jnp.array(translation).astype(float)

jaxsim/math/utils.py

@@ -5,8 +5,8 @@ import jaxsim.typing as jtp
 
 def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
     """
-    Provides a calculation for an array norm so that it is safe
-    to compute the gradient and handle NaNs.
+    Compute an array norm handling NaNs and making sure that
+    it is safe to get the gradient.
 
     Args:
         array: The array for which to compute the norm.

jaxsim/mujoco/loaders.py

@@ -22,7 +22,7 @@ def load_rod_model(
     model_name: str | None = None,
 ) -> rod.Model:
     """
-    Loads a ROD model from a URDF/SDF file or a ROD model.
+    Load a ROD model from a URDF/SDF file or a ROD model.
 
     Args:
         model_description: The URDF/SDF file or ROD model to load.
@@ -62,14 +62,16 @@ def load_rod_model(
 
 
 class RodModelToMjcf:
-    """"""
+    """
+    Class to convert a ROD model to a Mujoco MJCF string.
+    """
 
     @staticmethod
     def assets_from_rod_model(
         rod_model: rod.Model,
     ) -> dict[str, bytes]:
         """
-        Generates a dictionary of assets from a ROD model.
+        Generate a dictionary of assets from a ROD model.
 
         Args:
             rod_model: The ROD model to extract the assets from.
@@ -112,7 +114,7 @@ class RodModelToMjcf:
         floating_joint_name: str = "world_to_base",
     ) -> str:
         """
-        Adds a floating joint to a URDF string.
+        Add a floating joint to a URDF string.
 
         Args:
             urdf_string: The URDF string to modify.
@@ -171,7 +173,7 @@ class RodModelToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a ROD model to a Mujoco MJCF string.
+        Convert a ROD model to a Mujoco MJCF string.
 
         Args:
             rod_model: The ROD model to convert.
@@ -522,6 +524,10 @@ class RodModelToMjcf:
 
 
 class UrdfToMjcf:
+    """
+    Class to convert a URDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         urdf: str | pathlib.Path,
@@ -532,7 +538,7 @@ class UrdfToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a URDF file to a Mujoco MJCF string.
+        Convert a URDF file to a Mujoco MJCF string.
 
         Args:
             urdf: The URDF file to convert.
@@ -564,6 +570,10 @@ class UrdfToMjcf:
 
 
 class SdfToMjcf:
+    """
+    Class to convert a SDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         sdf: str | pathlib.Path,
@@ -574,7 +584,7 @@ class SdfToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a SDF file to a Mujoco MJCF string.
+        Convert a SDF file to a Mujoco MJCF string.
 
         Args:
             sdf: The SDF file to convert.