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

jaxsim/__init__.py

@@ -34,13 +34,6 @@ def _jnp_options() -> None:
         logging.info("Enabling JAX to use 64-bit precision")
         jax.config.update("jax_enable_x64", True)
 
-        import jax.numpy as jnp
-        import numpy as np
-
-        # Verify that 64-bit precision is correctly set.
-        if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
-            logging.warning("Failed to enable 64-bit precision in JAX")
-
     # Warn about experimental usage of 32-bit precision.
     else:
         logging.warning(

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev126'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev126')
+__version__ = version = '0.5.1.dev139'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev139')

jaxsim/api/com.py

@@ -279,7 +279,7 @@ def bias_acceleration(
         C_v̇_WL: jtp.Vector, C_v_WC: jtp.Vector, L_H_C: jtp.Matrix, L_v_LC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the body-fixed representation of the link bias acceleration
+        Convert the body-fixed representation of the link bias acceleration
         C_v̇_WL expressed in a generic frame C to the body-fixed representation L_v̇_WL.
         """
 

jaxsim/api/common.py

@@ -26,7 +26,7 @@ _R = TypeVar("_R")
 
 
 def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
-    """Applies a JAX named scope to a function for improved profiling and clarity."""
+    """Apply a JAX named scope to a function for improved profiling and clarity."""
 
     @functools.wraps(fn)
     def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:

jaxsim/api/contact.py

@@ -293,6 +293,9 @@ def in_contact(
 def estimate_good_soft_contacts_parameters(
     *args, **kwargs
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
+    """
+    Estimate good soft contacts parameters. Deprecated, use `estimate_good_contact_parameters` instead.
+    """
 
     msg = "This method is deprecated, please use `{}`."
     logging.warning(msg.format(estimate_good_contact_parameters.__name__))

jaxsim/api/data.py

@@ -456,7 +456,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     @jax.jit
     def generalized_velocity(self) -> jtp.Vector:
         r"""
-        Get the generalized velocity
+        Get the generalized velocity.
+
         :math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\, \boldsymbol{\omega}_{W,B};\, \mathbf{s}) \in \mathbb{R}^{6+n}`
 
         Returns:

jaxsim/api/kin_dyn_parameters.py

@@ -52,10 +52,16 @@ class KinDynParameters(JaxsimDataclass):
 
     @property
     def parent_array(self) -> jtp.Vector:
+        r"""
+        Return the parent array :math:`\lambda(i)` of the model.
+        """
         return self._parent_array.get()
 
     @property
     def support_body_array_bool(self) -> jtp.Matrix:
+        r"""
+        Return the boolean support parent array :math:`\kappa_{b}(i)` of the model.
+        """
         return self._support_body_array_bool.get()
 
     @staticmethod
@@ -648,7 +654,16 @@ class LinkParameters(JaxsimDataclass):
     def build_from_flat_parameters(
         index: jtp.IntLike, parameters: jtp.VectorLike
     ) -> LinkParameters:
+        """
+        Build a LinkParameters object from a flat vector of parameters.
+
+        Args:
+            index: The index of the link.
+            parameters: The flat vector of parameters.
 
+        Returns:
+            The LinkParameters object.
+        """
         index = jnp.array(index).squeeze().astype(int)
 
         m = jnp.array(parameters[0]).squeeze().astype(float)
@@ -772,7 +787,9 @@ class ContactParameters(JaxsimDataclass):
 
     @property
     def indices_of_enabled_collidable_points(self) -> npt.NDArray:
-
+        """
+        Return the indices of the enabled collidable points.
+        """
         return np.where(np.array(self.enabled))[0]
 
     @staticmethod

jaxsim/api/model.py

@@ -63,6 +63,9 @@ class JaxSimModel(JaxsimDataclass):
 
     @property
     def description(self) -> ModelDescription:
+        """
+        Return the model description.
+        """
         return self._description.get()
 
     def __eq__(self, other: JaxSimModel) -> bool:
@@ -1015,7 +1018,7 @@ def forward_dynamics_aba(
         W_v̇_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WB: jtp.Vector, W_v_WC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the inertial-fixed apparent base acceleration W_v̇_WB to
+        Convert the inertial-fixed apparent base acceleration W_v̇_WB to
         another representation C_v̇_WB expressed in a generic frame C.
         """
 
@@ -1376,7 +1379,7 @@ def inverse_dynamics(
 
     def to_inertial(C_v̇_WB, W_H_C, C_v_WB, W_v_WC):
         """
-        Helper to convert the active representation of the base acceleration C_v̇_WB
+        Convert the active representation of the base acceleration C_v̇_WB
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
@@ -1825,7 +1828,7 @@ def link_bias_accelerations(
         C_v̇_WB: jtp.Vector, C_v_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the active representation of the base acceleration C_v̇_WB
+        Convert the active representation of the base acceleration C_v̇_WB
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
@@ -1961,7 +1964,7 @@ def link_bias_accelerations(
         L_v̇_WL: jtp.Vector, L_v_WL: jtp.Vector, C_H_L: jtp.Matrix, L_v_CL: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the body-fixed apparent acceleration L_v̇_WL to
+        Convert the body-fixed apparent acceleration L_v̇_WL to
         another representation C_v̇_WL expressed in a generic frame C.
         """
 

jaxsim/api/ode.py

@@ -15,12 +15,32 @@ from .ode_data import ODEState
 
 
 class SystemDynamicsFromModelAndData(Protocol):
+    """
+    Protocol defining the signature of a function computing the system dynamics
+    given a model and data object.
+    """
+
     def __call__(
         self,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         **kwargs: dict[str, Any],
-    ) -> tuple[ODEState, dict[str, Any]]: ...
+    ) -> tuple[ODEState, dict[str, Any]]:
+        """
+        Compute the system dynamics given a model and data object.
+
+        Args:
+            model: The model to consider.
+            data: The data of the considered model.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            A tuple with an `ODEState` object storing in each of its attributes the
+            corresponding derivative, and the dictionary of auxiliary data returned
+            by the system dynamics evaluation.
+        """
+
+        pass
 
 
 def wrap_system_dynamics_for_integration(

jaxsim/exceptions.py

@@ -17,6 +17,8 @@ def raise_if(
         msg:
             The message to display when the exception is raised. The message can be a
             format string (fmt), whose fields are filled with the args and kwargs.
+        *args: The arguments to fill the format string.
+        **kwargs: The keyword arguments to fill the format string
     """
 
     # Disable host callback if running on unsupported hardware or if the user
@@ -61,6 +63,9 @@ def raise_if(
 def raise_runtime_error_if(
     condition: bool | jax.Array, msg: str, *args, **kwargs
 ) -> None:
+    """
+    Raise a RuntimeError if a condition is met. Useful in jit-compiled functions.
+    """
 
     return raise_if(condition, RuntimeError, msg, *args, **kwargs)
 
@@ -68,5 +73,8 @@ def raise_runtime_error_if(
 def raise_value_error_if(
     condition: bool | jax.Array, msg: str, *args, **kwargs
 ) -> None:
+    """
+    Raise a ValueError if a condition is met. Useful in jit-compiled functions.
+    """
 
     return raise_if(condition, ValueError, msg, *args, **kwargs)

jaxsim/integrators/common.py

@@ -36,9 +36,25 @@ PyTreeType = TypeVar("PyTreeType", bound=jtp.PyTree)
 
 
 class SystemDynamics(Protocol[State, StateDerivative]):
+    """
+    Protocol defining the system dynamics.
+    """
+
     def __call__(
         self, x: State, t: Time, **kwargs
-    ) -> tuple[StateDerivative, dict[str, Any]]: ...
+    ) -> tuple[StateDerivative, dict[str, Any]]:
+        """
+        Compute the state derivative of the system.
+
+        Args:
+            x: The state of the system.
+            t: The time of the system.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The state derivative of the system and the auxiliary dictionary.
+        """
+        pass
 
 
 # =======================
@@ -48,6 +64,9 @@ class SystemDynamics(Protocol[State, StateDerivative]):
 
 @jax_dataclasses.pytree_dataclass
 class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
+    """
+    Factory class for integrators.
+    """
 
     dynamics: Static[SystemDynamics[State, StateDerivative]] = dataclasses.field(
         repr=False, hash=False, compare=False, kw_only=True
@@ -110,6 +129,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Perform a single integration step.
+        """
         pass
 
     def init(
@@ -121,6 +143,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         include_dynamics_aux_dict: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
+        """
+        Initialize the integrator. This method is deprecated.
+        """
 
         logging.warning(
             "The 'init' method has been deprecated. There is no need to call it."
@@ -131,16 +156,28 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
 @jax_dataclasses.pytree_dataclass
 class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]):
+    """
+    Base class for explicit Runge-Kutta integrators.
+
+    Attributes:
+        A: The Runge-Kutta matrix.
+        b: The weights coefficients.
+        c: The nodes coefficients.
+        order_of_bT_rows: The order of the solution.
+        row_index_of_solution: The row of the integration output corresponding to the final solution.
+        fsal_enabled_if_supported: Whether to enable the FSAL property, if supported.
+        index_of_fsal: The index of the intermediate derivative to be used as the first derivative of the next iteration.
+    """
 
     # The Runge-Kutta matrix.
-    A: ClassVar[jtp.Matrix]
+    A: jtp.Matrix
 
     # The weights coefficients.
     # Note that in practice we typically use its transpose `b.transpose()`.
-    b: ClassVar[jtp.Matrix]
+    b: jtp.Matrix
 
     # The nodes coefficients.
-    c: ClassVar[jtp.Vector]
+    c: jtp.Vector
 
     # Define the order of the solution.
     # It should have as many elements as the number of rows of `b.transpose()`.
@@ -156,10 +193,16 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     @property
     def has_fsal(self) -> bool:
+        """
+        Check if the integrator supports the FSAL property.
+        """
         return self.fsal_enabled_if_supported and self.index_of_fsal is not None
 
     @property
     def order(self) -> int:
+        """
+        Return the order of the integrator.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution]
 
     @override
@@ -183,28 +226,31 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         Returns:
             The integrator object.
         """
+        A = cls.__dataclass_fields__["A"].default_factory()
+        b = cls.__dataclass_fields__["b"].default_factory()
+        c = cls.__dataclass_fields__["c"].default_factory()
 
         # Check validity of the Butcher tableau.
-        if not ExplicitRungeKutta.butcher_tableau_is_valid(A=cls.A, b=cls.b, c=cls.c):
+        if not ExplicitRungeKutta.butcher_tableau_is_valid(A=A, b=b, c=c):
             raise ValueError("The Butcher tableau of this class is not valid.")
 
         # Check that b.T has enough rows based on the configured index of the solution.
-        if cls.row_index_of_solution >= cls.b.T.shape[0]:
+        if cls.row_index_of_solution >= b.T.shape[0]:
             msg = "The index of the solution ({}-th row of `b.T`) is out of range ({})."
-            raise ValueError(msg.format(cls.row_index_of_solution, cls.b.T.shape[0]))
+            raise ValueError(msg.format(cls.row_index_of_solution, b.T.shape[0]))
 
         # Check that the tuple containing the order of the b.T rows matches the number
         # of the b.T rows.
-        if len(cls.order_of_bT_rows) != cls.b.T.shape[0]:
+        if len(cls.order_of_bT_rows) != b.T.shape[0]:
             msg = "Wrong size of 'order_of_bT_rows' ({}), should be {}."
-            raise ValueError(msg.format(len(cls.order_of_bT_rows), cls.b.T.shape[0]))
+            raise ValueError(msg.format(len(cls.order_of_bT_rows), b.T.shape[0]))
 
         # Check if the Butcher tableau supports FSAL (first-same-as-last).
         # If it does, store the index of the intermediate derivative to be used as the
         # first derivative of the next iteration.
         has_fsal, index_of_fsal = (  # noqa: F841
             ExplicitRungeKutta.butcher_tableau_supports_fsal(
-                A=cls.A, b=cls.b, c=cls.c, index_of_solution=cls.row_index_of_solution
+                A=A, b=b, c=c, index_of_solution=cls.row_index_of_solution
             )
         )
 
@@ -221,6 +267,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Perform a single integration step.
+        """
 
         # Here z is a batched state with as many batch elements as b.T rows.
         # Note that z has multiple batches only if b.T has more than one row,
@@ -331,7 +380,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         def scan_body(
             carry: jax.Array, i: int | jax.Array
         ) -> tuple[jax.Array, dict[str, Any]]:
-            """"""
+            """
+            Compute the kᵢ derivative of the Runge-Kutta stage.
+            """
 
             # Unpack the carry, i.e. the stacked kᵢ vectors.
             K = carry
@@ -498,6 +549,16 @@ class ExplicitRungeKuttaSO3Mixin:
     def post_process_state(
         cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
     ) -> js.ode_data.ODEState:
+        r"""
+        Post-process the integrated state at :math:`t_f = t_0 + \Delta t` so that the
+        quaternion is normalized.
+
+        Args:
+            x0: The initial state of the system.
+            t0: The initial time of the system.
+            xf: The final state of the system obtain through the integration.
+            dt: The time step used for the integration.
+        """
 
         # Extract the initial base quaternion.
         W_Q_B_t0 = x0.physics_model.base_quaternion

jaxsim/integrators/fixed_step.py

@@ -1,3 +1,4 @@
+import dataclasses
 from typing import ClassVar, Generic
 
 import jax.numpy as jnp
@@ -17,69 +18,107 @@ ODEStateDerivative = js.ode_data.ODEState
 
 @jax_dataclasses.pytree_dataclass
 class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Forward Euler integrator.
+    """
 
-    A: ClassVar[jtp.Matrix] = jnp.atleast_2d(0).astype(float)
-
-    b: ClassVar[jtp.Matrix] = jnp.atleast_2d(1).astype(float).transpose()
+    A: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.atleast_2d(0).astype(float), compare=False
+    )
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.atleast_2d(1).astype(float), compare=False
+    )
 
-    c: ClassVar[jtp.Vector] = jnp.atleast_1d(0).astype(float)
+    c: jtp.Vector = dataclasses.field(
+        default_factory=lambda: jnp.atleast_1d(0).astype(float), compare=False
+    )
 
-    row_index_of_solution: ClassVar[int] = 0
-    order_of_bT_rows: ClassVar[tuple[int, ...]] = (1,)
+    row_index_of_solution: int = 0
+    order_of_bT_rows: tuple[int, ...] = (1,)
+    index_of_fsal: jtp.IntLike | None = None
+    fsal_enabled_if_supported: bool = False
 
 
 @jax_dataclasses.pytree_dataclass
 class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Heun's second-order integrator.
+    """
+
+    A: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [
+                [0, 0],
+                [1, 0],
+            ]
+        ).astype(float),
+        compare=False,
+    )
 
-    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]),
-        )
-        .astype(float)
-        .transpose()
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: (
+            jnp.atleast_2d(
+                jnp.array([1 / 2, 1 / 2]),
+            )
+            .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
     order_of_bT_rows: ClassVar[tuple[int, ...]] = (2,)
+    index_of_fsal: jtp.IntLike | None = None
+    fsal_enabled_if_supported: bool = False
 
 
 @jax_dataclasses.pytree_dataclass
 class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Fourth-order Runge-Kutta integrator.
+    """
+
+    A: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [
+                [0, 0, 0, 0],
+                [1 / 2, 0, 0, 0],
+                [0, 1 / 2, 0, 0],
+                [0, 0, 1, 0],
+            ]
+        ).astype(float),
+        compare=False,
+    )
 
-    A: ClassVar[jtp.Matrix] = jnp.array(
-        [
-            [0, 0, 0, 0],
-            [1 / 2, 0, 0, 0],
-            [0, 1 / 2, 0, 0],
-            [0, 0, 1, 0],
-        ]
-    ).astype(float)
-
-    b: ClassVar[jtp.Matrix] = (
-        jnp.atleast_2d(
-            jnp.array([1 / 6, 1 / 3, 1 / 3, 1 / 6]),
-        )
-        .astype(float)
-        .transpose()
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: (
+            jnp.atleast_2d(
+                jnp.array([1 / 6, 1 / 3, 1 / 3, 1 / 6]),
+            )
+            .astype(float)
+            .transpose()
+        ),
+        compare=False,
     )
 
-    c: ClassVar[jtp.Vector] = jnp.array(
-        [0, 1 / 2, 1 / 2, 1],
-    ).astype(float)
+    c: jtp.Vector = dataclasses.field(
+        default_factory=lambda: jnp.array(
+            [0, 1 / 2, 1 / 2, 1],
+        ).astype(float),
+        compare=False,
+    )
 
     row_index_of_solution: ClassVar[int] = 0
     order_of_bT_rows: ClassVar[tuple[int, ...]] = (4,)
+    index_of_fsal: jtp.IntLike | None = None
+    fsal_enabled_if_supported: bool = False
 
 
 # ===============================================================================
@@ -89,14 +128,26 @@ class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class ForwardEulerSO3(ExplicitRungeKuttaSO3Mixin, ForwardEuler[js.ode_data.ODEState]):
+    """
+    Forward Euler integrator for SO(3) states.
+    """
+
     pass
 
 
 @jax_dataclasses.pytree_dataclass
 class Heun2SO3(ExplicitRungeKuttaSO3Mixin, Heun2[js.ode_data.ODEState]):
+    """
+    Heun's second-order integrator for SO(3) states.
+    """
+
     pass
 
 
 @jax_dataclasses.pytree_dataclass
 class RungeKutta4SO3(ExplicitRungeKuttaSO3Mixin, RungeKutta4[js.ode_data.ODEState]):
+    """
+    Fourth-order Runge-Kutta integrator for SO(3) states.
+    """
+
     pass