jaxsim 0.5.1.dev133__tar.gz → 0.5.1.dev143__tar.gz

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev133
+Version: 0.5.1.dev143
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/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-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/_version.py

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

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/api/model.py

@@ -2295,9 +2295,6 @@ def step(
         # Hence, here we need to reset the velocity after each impact to guarantee that
         # the linear velocity of the active collidable points is zero.
         case jaxsim.rbda.contacts.RigidContacts():
-            assert isinstance(
-                data_tf.contacts_params, jaxsim.rbda.contacts.RigidContactsParams
-            )
 
             # Raise runtime error for not supported case in which Rigid contacts and
             # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
@@ -2331,12 +2328,13 @@ def step(
                     indices_of_enabled_collidable_points
                 ]
                 M = js.model.free_floating_mass_matrix(model, data_tf)
+                BW_ν_pre_impact = data_tf.generalized_velocity()
 
                 # Compute the impact velocity.
                 # It may be discontinuous in case new contacts are made.
-                BW_nu_post_impact = (
+                BW_ν_post_impact = (
                     jaxsim.rbda.contacts.RigidContacts.compute_impact_velocity(
-                        data=data_tf,
+                        generalized_velocity=BW_ν_pre_impact,
                         inactive_collidable_points=(δ <= 0),
                         M=M,
                         J_WC=J_WC,
@@ -2344,8 +2342,8 @@ def step(
                 )
 
                 # Reset the generalized velocity.
-                data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
-                data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+                data_tf = data_tf.reset_base_velocity(BW_ν_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_ν_post_impact[6:])
 
     # Restore the input velocity representation.
     data_tf = data_tf.replace(

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/integrators/common.py

@@ -170,14 +170,14 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
     """
 
     # 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()`.
@@ -226,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
             )
         )
 

jaxsim-0.5.1.dev143/src/jaxsim/integrators/fixed_step.py

@@ -0,0 +1,153 @@
+import dataclasses
+from typing import ClassVar, Generic
+
+import jax.numpy as jnp
+import jax_dataclasses
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+
+from .common import ExplicitRungeKutta, ExplicitRungeKuttaSO3Mixin, PyTreeType
+
+ODEStateDerivative = js.ode_data.ODEState
+
+# =====================================================
+# Explicit Runge-Kutta integrators operating on PyTrees
+# =====================================================
+
+
+@jax_dataclasses.pytree_dataclass
+class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Forward Euler integrator.
+    """
+
+    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: jtp.Vector = dataclasses.field(
+        default_factory=lambda: jnp.atleast_1d(0).astype(float), compare=False
+    )
+
+    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,
+    )
+
+    b: jtp.Matrix = dataclasses.field(
+        default_factory=lambda: (
+            jnp.atleast_2d(
+                jnp.array([1 / 2, 1 / 2]),
+            )
+            .astype(float)
+            .transpose()
+        ),
+        compare=False,
+    )
+
+    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,
+    )
+
+    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: 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
+
+
+# ===============================================================================
+# Explicit Runge-Kutta integrators operating on ODEState and integrating on SO(3)
+# ===============================================================================
+
+
+@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

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/integrators/variable_step.py

@@ -1,3 +1,4 @@
+import dataclasses
 import functools
 from typing import Any, ClassVar, Generic
 
@@ -254,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,
@@ -573,16 +577,18 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             **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(
@@ -611,35 +617,47 @@ 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,
     )
 
-    c: ClassVar[jtp.Vector] = jnp.array(
-        [0, 1],
-    ).astype(float)
+    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: 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):
@@ -647,31 +665,40 @@ class BogackiShampineSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mi
     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-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/math/adjoint.py

@@ -13,8 +13,8 @@ class Adjoint:
 
     @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:
@@ -22,14 +22,17 @@ class Adjoint:
         Create an adjoint matrix from a quaternion and a translation.
 
         Args:
-            quaternion: A quaternion vector (4D) representing orientation.
-            translation: A translation vector (3D).
-            inverse: Whether to compute the inverse adjoint.
-            normalize_quaternion: Whether to normalize the quaternion before creating the adjoint.
+            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.
 
         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
 
@@ -64,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: A 3x3 rotation matrix.
-            translation: A translation vector (3D).
-            inverse: Whether to compute the inverse adjoint. Default is False.
+            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
 

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/math/transform.py

@@ -11,8 +11,8 @@ class Transform:
 
     @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:
@@ -30,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)
 
@@ -47,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:
         """
@@ -62,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-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim/rbda/contacts/rigid.py

@@ -173,7 +173,7 @@ class RigidContacts(ContactModel):
         inactive_collidable_points: jtp.ArrayLike,
         M: jtp.MatrixLike,
         J_WC: jtp.MatrixLike,
-        data: js.data.JaxSimModelData,
+        generalized_velocity: jtp.VectorLike,
     ) -> jtp.Vector:
         """
         Return the new velocity of the system after a potential impact.
@@ -182,55 +182,37 @@ class RigidContacts(ContactModel):
             inactive_collidable_points: The activation state of the collidable points.
             M: The mass matrix of the system (in mixed representation).
             J_WC: The Jacobian matrix of the collidable points (in mixed representation).
-            data: The `JaxSimModelData` instance.
-        """
+            generalized_velocity: The generalized velocity of the system.
 
-        def impact_velocity(
-            inactive_collidable_points: jtp.ArrayLike,
-            nu_pre: jtp.ArrayLike,
-            M: jtp.MatrixLike,
-            J_WC: jtp.MatrixLike,
-            data: js.data.JaxSimModelData,
-        ):
-            # Compute system velocity after impact maintaining zero linear velocity of active points
-            with data.switch_velocity_representation(VelRepr.Mixed):
-                sl = jnp.s_[:, 0:3, :]
-                Jl_WC = J_WC[sl]
-                # Zero out the jacobian rows of inactive points
-                Jl_WC = jnp.vstack(
-                    jnp.where(
-                        inactive_collidable_points[:, jnp.newaxis, jnp.newaxis],
-                        jnp.zeros_like(Jl_WC),
-                        Jl_WC,
-                    )
-                )
+        Note:
+            The mass matrix `M`, the Jacobian `J_WC`, and the generalized velocity `generalized_velocity`
+            must be expressed in the same velocity representation.
+        """
 
-                A = jnp.vstack(
-                    [
-                        jnp.hstack([M, -Jl_WC.T]),
-                        jnp.hstack(
-                            [Jl_WC, jnp.zeros((Jl_WC.shape[0], Jl_WC.shape[0]))]
-                        ),
-                    ]
-                )
-                b = jnp.hstack([M @ nu_pre, jnp.zeros(Jl_WC.shape[0])])
-                x = jnp.linalg.lstsq(A, b)[0]
-                nu_post = x[0 : M.shape[0]]
+        # Compute system velocity after impact maintaining zero linear velocity of active points.
+        sl = jnp.s_[:, 0:3, :]
+        Jl_WC = J_WC[sl]
+
+        # Zero out the jacobian rows of inactive points.
+        Jl_WC = jnp.vstack(
+            jnp.where(
+                inactive_collidable_points[:, jnp.newaxis, jnp.newaxis],
+                jnp.zeros_like(Jl_WC),
+                Jl_WC,
+            )
+        )
 
-                return nu_post
+        A = jnp.vstack(
+            [
+                jnp.hstack([M, -Jl_WC.T]),
+                jnp.hstack([Jl_WC, jnp.zeros((Jl_WC.shape[0], Jl_WC.shape[0]))]),
+            ]
+        )
+        b = jnp.hstack([M @ generalized_velocity, jnp.zeros(Jl_WC.shape[0])])
 
-        with data.switch_velocity_representation(VelRepr.Mixed):
-            BW_ν_pre_impact = data.generalized_velocity()
-
-            BW_ν_post_impact = impact_velocity(
-                data=data,
-                inactive_collidable_points=inactive_collidable_points,
-                nu_pre=BW_ν_pre_impact,
-                M=M,
-                J_WC=J_WC,
-            )
+        BW_ν_post_impact = jnp.linalg.lstsq(A, b)[0]
 
-        return BW_ν_post_impact
+        return BW_ν_post_impact[0 : M.shape[0]]
 
     @jax.jit
     def compute_contact_forces(

{jaxsim-0.5.1.dev133 → jaxsim-0.5.1.dev143}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev133
+Version: 0.5.1.dev143
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>