jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/integrators/common.py

@@ -1,13 +1,16 @@
 import abc
 import dataclasses
-from typing import Any, ClassVar, Generic, Protocol, Type, TypeVar
+from typing import Any, ClassVar, Generic, Protocol, TypeVar
 
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
 from jax_dataclasses import Static
 
+import jaxsim.api as js
+import jaxsim.math
 import jaxsim.typing as jtp
+from jaxsim import exceptions, logging
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -25,17 +28,33 @@ except ImportError:
 # Generic types
 # =============
 
-Time = jax.typing.ArrayLike
-TimeStep = jax.typing.ArrayLike
+Time = jtp.FloatLike
+TimeStep = jtp.FloatLike
 State = NextState = TypeVar("State")
 StateDerivative = TypeVar("StateDerivative")
 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
 
 
 # =======================
@@ -45,20 +64,20 @@ class SystemDynamics(Protocol[State, StateDerivative]):
 
 @jax_dataclasses.pytree_dataclass
 class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
-
-    AuxDictDynamicsKey: ClassVar[str] = "aux_dict_dynamics"
+    """
+    Factory class for integrators.
+    """
 
     dynamics: Static[SystemDynamics[State, StateDerivative]] = dataclasses.field(
         repr=False, hash=False, compare=False, kw_only=True
     )
 
-    params: dict[str, Any] = dataclasses.field(
-        default_factory=dict, repr=False, hash=False, compare=False, kw_only=True
-    )
-
     @classmethod
     def build(
-        cls: Type[Self], *, dynamics: SystemDynamics[State, StateDerivative], **kwargs
+        cls: type[Self],
+        *,
+        dynamics: SystemDynamics[State, StateDerivative],
+        **kwargs,
     ) -> Self:
         """
         Build the integrator object.
@@ -71,7 +90,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             The integrator object.
         """
 
-        return cls(dynamics=dynamics, **kwargs)  # noqa
+        return cls(dynamics=dynamics, **kwargs)
 
     def step(
         self,
@@ -79,9 +98,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         t0: Time,
         dt: TimeStep,
         *,
-        params: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
         **kwargs,
-    ) -> tuple[State, dict[str, Any]]:
+    ) -> tuple[NextState, dict[str, Any]]:
         """
         Perform a single integration step.
 
@@ -89,25 +108,30 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             x0: The initial state of the system.
             t0: The initial time of the system.
             dt: The time step of the integration.
-            params: The auxiliary dictionary of the integrator.
+            metadata: The state auxiliary dictionary of the integrator.
             **kwargs: Additional keyword arguments.
 
         Returns:
             The final state of the system and the updated auxiliary dictionary.
         """
 
-        with self.editable(validate=False) as integrator:
-            integrator.params = params
-
-        with integrator.mutable_context(mutability=Mutability.MUTABLE):
-            xf = integrator(x0, t0, dt, **kwargs)
+        metadata = metadata if metadata is not None else {}
 
-        assert Integrator.AuxDictDynamicsKey in integrator.params
+        with self.mutable_context(mutability=Mutability.MUTABLE) as integrator:
+            xf, metadata_step = integrator(x0, t0, dt, **kwargs)
 
-        return xf, integrator.params
+        return (
+            xf,
+            metadata | metadata_step,
+        )
 
     @abc.abstractmethod
-    def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
+    def __call__(
+        self, x0: State, t0: Time, dt: TimeStep, **kwargs
+    ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Perform a single integration step.
+        """
         pass
 
     def init(
@@ -116,56 +140,44 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         t0: Time,
         dt: TimeStep,
         *,
-        key: jax.Array | None = None,
+        include_dynamics_aux_dict: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
         """
-        Initialize the integrator.
-
-        Args:
-            x0: The initial state of the system.
-            t0: The initial time of the system.
-            dt: The time step of the integration.
-            key: An optional random key to initialize the integrator.
-
-        Returns:
-            The auxiliary dictionary of the integrator.
-
-        Note:
-            This method should have the same signature as the inherited `__call__`
-            method, including additional kwargs.
-
-        Note:
-            If the integrator supports FSAL, the pair `(x0, t0)` must match the real
-            initial state and time of the system, otherwise the initial derivative of
-            the first step will be wrong.
+        Initialize the integrator. This method is deprecated.
         """
 
-        _, aux_dict_dynamics = self.dynamics(x0, t0)
-
-        with self.editable(validate=False) as integrator:
-            _ = integrator(x0, t0, dt, **kwargs)
-            aux_dict_step = integrator.params
-
-        if Integrator.AuxDictDynamicsKey in aux_dict_dynamics:
-            msg = "You cannot create a key '{}' in the __call__ method."
-            raise KeyError(msg.format(Integrator.AuxDictDynamicsKey))
+        logging.warning(
+            "The 'init' method has been deprecated. There is no need to call it."
+        )
 
-        return {Integrator.AuxDictDynamicsKey: aux_dict_dynamics} | aux_dict_step
+        return {}
 
 
 @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[jax.typing.ArrayLike]
+    A: jtp.Matrix
 
     # The weights coefficients.
     # Note that in practice we typically use its transpose `b.transpose()`.
-    b: ClassVar[jax.typing.ArrayLike]
+    b: jtp.Matrix
 
     # The nodes coefficients.
-    c: ClassVar[jax.typing.ArrayLike]
+    c: jtp.Vector
 
     # Define the order of the solution.
     # It should have as many elements as the number of rows of `b.transpose()`.
@@ -181,16 +193,22 @@ 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
     @classmethod
     def build(
-        cls: Type[Self],
+        cls: type[Self],
         *,
         dynamics: SystemDynamics[State, StateDerivative],
         fsal_enabled_if_supported: jtp.BoolLike = True,
@@ -208,37 +226,32 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         Returns:
             The integrator object.
         """
-
-        # Adjust the shape of the tableau coefficients.
-        c = jnp.atleast_1d(cls.c.squeeze())
-        b = jnp.atleast_2d(jnp.vstack(cls.b.squeeze()))
-        A = jnp.atleast_2d(cls.A.squeeze())
+        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=A, b=b, c=c):
             raise ValueError("The Butcher tableau of this class is not valid.")
 
-        # Store the adjusted shapes of the tableau coefficients.
-        cls.c = c
-        cls.b = b
-        cls.A = A
-
         # 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 = ExplicitRungeKutta.butcher_tableau_supports_fsal(
-            A=cls.A, b=cls.b, c=cls.c, index_of_solution=cls.row_index_of_solution
+        has_fsal, index_of_fsal = (  # noqa: F841
+            ExplicitRungeKutta.butcher_tableau_supports_fsal(
+                A=A, b=b, c=c, index_of_solution=cls.row_index_of_solution
+            )
         )
 
         # Build the integrator object.
@@ -251,15 +264,22 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         return integrator
 
-    def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
+    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,
         # e.g. in Butcher tableau of embedded schemes.
-        z = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
+        z, aux_dict = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
 
         # The next state is the batch element located at the configured index of solution.
-        return jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
+        next_state = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
+
+        return next_state, aux_dict
 
     @classmethod
     def integrate_rk_stage(
@@ -294,13 +314,13 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         """
 
         op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
-        return jax.tree_util.tree_map(op, x0, k)
+        return jax.tree.map(op, x0, k)
 
     @classmethod
     def post_process_state(
         cls, x0: State, t0: Time, xf: NextState, dt: TimeStep
     ) -> NextState:
-        """
+        r"""
         Post-process the integrated state at :math:`t_f = t_0 + \Delta t`.
 
         Args:
@@ -317,7 +337,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     def _compute_next_state(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
-    ) -> NextState:
+    ) -> tuple[NextState, dict[str, Any]]:
         """
         Compute the next state of the system, returning all the output states.
 
@@ -337,33 +357,42 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         b = self.b
         A = self.A
 
+        # Extract metadata from the kwargs.
+        metadata = kwargs.pop("metadata", {})
+
         # Close f over optional kwargs.
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_util.tree_map(
-            lambda l: jnp.repeat(jnp.zeros_like(l)[jnp.newaxis, ...], c.size, axis=0),
-            x0,
+        carry0 = jax.tree.map(
+            lambda l: jnp.zeros((c.size, *l.shape), dtype=l.dtype), x0
         )
 
-        # Apply FSAL property by passing ẋ0 = f(x0, t0) from the previous iteration.
-        get_ẋ0 = lambda: self.params.get("dxdt0", f(x0, t0)[0])
+        # Closure on metadata to either evaluate the dynamics at the initial state
+        # or to use the previous state derivative (only integrators supporting FSAL).
+        def get_ẋ0_and_aux_dict() -> tuple[StateDerivative, dict[str, Any]]:
+            ẋ0, aux_dict = f(x0, t0)
+            return metadata.get("dxdt0", ẋ0), aux_dict
 
         # We use a `jax.lax.scan` to compile the `f` function only once.
         # Otherwise, if we compute e.g. for RK4 sequentially, the jit-compiled code
         # would include 4 repetitions of the `f` logic, making everything extremely slow.
-        def scan_body(carry: jax.Array, i: int | jax.Array) -> tuple[jax.Array, None]:
-            """"""
+        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
 
             # Define the computation of the Runge-Kutta stage.
-            def compute_ki() -> jax.Array:
+            def compute_ki() -> tuple[jax.Array, dict[str, Any]]:
 
-                # Compute ∑ⱼ aᵢⱼ kⱼ
+                # Compute ∑ⱼ aᵢⱼ kⱼ.
                 op_sum_ak = lambda k: jnp.einsum("s,s...->...", A[i], k)
-                sum_ak = jax.tree_util.tree_map(op_sum_ak, K)
+                sum_ak = jax.tree.map(op_sum_ak, K)
 
                 # Compute the next state for the kᵢ evaluation.
                 # Note that this is not a Δt integration since aᵢⱼ could be fractional.
@@ -372,25 +401,26 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
                 # Compute the next time for the kᵢ evaluation.
                 ti = t0 + c[i] * Δt
 
-                # This is kᵢ = f(xᵢ, tᵢ).
-                return f(xi, ti)[0]
+                # Evaluate the dynamics.
+                ki, aux_dict = f(xi, ti)
+                return ki, aux_dict
 
             # This selector enables FSAL property in the first iteration (i=0).
-            ki = jax.lax.cond(
+            ki, aux_dict = jax.lax.cond(
                 pred=jnp.logical_and(i == 0, self.has_fsal),
-                true_fun=get_ẋ0,
+                true_fun=get_ẋ0_and_aux_dict,
                 false_fun=compute_ki,
             )
 
             # Store the kᵢ derivative in K.
             op = lambda l_k, l_ki: l_k.at[i].set(l_ki)
-            K = jax.tree_util.tree_map(op, K, ki)
+            K = jax.tree.map(op, K, ki)
 
             carry = K
-            return carry, None
+            return carry, aux_dict
 
         # Compute the state derivatives kᵢ.
-        K, _ = jax.lax.scan(
+        K, aux_dict = jax.lax.scan(
             f=scan_body,
             init=carry0,
             xs=jnp.arange(c.size),
@@ -398,12 +428,13 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_map(lambda l: l[self.index_of_fsal], K)
+            # Store the first derivative of the next step in the metadata.
+            metadata["dxdt0"] = jax.tree.map(lambda l: l[self.index_of_fsal], K)
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.
         op = lambda x0, k: x0 + Δt * jnp.einsum("zs,s...->z...", b.T, k)
-        z = jax.tree_util.tree_map(op, x0, K)
+        z = jax.tree.map(op, x0, K)
 
         # Transform the final state of the integration.
         # This allows to inject custom logic, if needed.
@@ -411,11 +442,11 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             lambda xf: self.post_process_state(x0=x0, t0=t0, xf=xf, dt=dt)
         )(z)
 
-        return z_transformed
+        return z_transformed, aux_dict | {"metadata": metadata}
 
     @staticmethod
     def butcher_tableau_is_valid(
-        A: jax.typing.ArrayLike, b: jax.typing.ArrayLike, c: jax.typing.ArrayLike
+        A: jtp.Matrix, b: jtp.Matrix, c: jtp.Vector
     ) -> jtp.Bool:
         """
         Check if the Butcher tableau is valid.
@@ -441,7 +472,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         return valid
 
     @staticmethod
-    def butcher_tableau_is_explicit(A: jax.typing.ArrayLike) -> jtp.Bool:
+    def butcher_tableau_is_explicit(A: jtp.Matrix) -> jtp.Bool:
         """
         Check if the Butcher tableau corresponds to an explicit integration scheme.
 
@@ -456,11 +487,11 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     @staticmethod
     def butcher_tableau_supports_fsal(
-        A: jax.typing.ArrayLike,
-        b: jax.typing.ArrayLike,
-        c: jax.typing.ArrayLike,
+        A: jtp.Matrix,
+        b: jtp.Matrix,
+        c: jtp.Vector,
         index_of_solution: jtp.IntLike = 0,
-    ) -> [bool, int | None]:
+    ) -> tuple[bool, int | None]:
         """
         Check if the Butcher tableau supports the FSAL (first-same-as-last) property.
 
@@ -481,7 +512,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             raise ValueError("The Butcher tableau is not valid.")
 
         if not ExplicitRungeKutta.butcher_tableau_is_explicit(A=A):
-            return False
+            return False, None
 
         if index_of_solution >= b.T.shape[0]:
             msg = "The index of the solution (i-th row of `b.T`) is out of range."
@@ -505,4 +536,57 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         # Return the index of the row of A providing the fsal derivative (that is the
         # possibly intermediate kᵢ derivative).
         # Note that if multiple rows match (it should not), we return the first match.
-        return True, int(jnp.where(rows_of_A_with_fsal == True)[0].tolist()[0])
+        return True, int(jnp.where(rows_of_A_with_fsal)[0].tolist()[0])
+
+
+class ExplicitRungeKuttaSO3Mixin:
+    """
+    Mixin class to apply over explicit RK integrators defined on
+    `PyTreeType = ODEState` to integrate the quaternion on SO(3).
+    """
+
+    @classmethod
+    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
+
+        # We assume that the initial quaternion is already unary.
+        exceptions.raise_runtime_error_if(
+            condition=~jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0),
+            msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
+        )
+
+        # Get the angular velocity ω to integrate the quaternion.
+        # This velocity ω[t0] is computed in the previous timestep by averaging the kᵢ
+        # corresponding to the active RK-based scheme. Therefore, by using the ω[t0],
+        # we obtain an explicit RK scheme operating on the SO(3) manifold.
+        # Note that the current integrator is not a semi-implicit scheme, therefore
+        # using the final ω[tf] would be not correct.
+        W_ω_WB_t0 = x0.physics_model.base_angular_velocity
+
+        # Integrate the quaternion on SO(3).
+        W_Q_B_tf = jaxsim.math.Quaternion.integration(
+            quaternion=W_Q_B_t0,
+            dt=dt,
+            omega=W_ω_WB_t0,
+            omega_in_body_fixed=False,
+        )
+
+        # Replace the quaternion in the final state.
+        return xf.replace(
+            physics_model=xf.physics_model.replace(base_quaternion=W_Q_B_tf),
+            validate=True,
+        )