jaxsim 0.4.1.dev24__py3-none-any.whl → 0.4.2__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.1.dev24'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev24')
+__version__ = version = '0.4.2'
+__version_tuple__ = version_tuple = (0, 4, 2)

jaxsim/api/com.py

@@ -137,9 +137,9 @@ def centroidal_momentum_jacobian(
 
     match data.velocity_representation:
         case VelRepr.Inertial | VelRepr.Mixed:
-            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)  # noqa: F841
         case VelRepr.Body:
-            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
         case _:
             raise ValueError(data.velocity_representation)
 
@@ -172,9 +172,9 @@ def locked_centroidal_spatial_inertia(
 
     match data.velocity_representation:
         case VelRepr.Inertial | VelRepr.Mixed:
-            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)  # noqa: F841
         case VelRepr.Body:
-            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
         case _:
             raise ValueError(data.velocity_representation)
 
@@ -290,14 +290,14 @@ def bias_acceleration(
 
         case VelRepr.Inertial:
 
-            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL
-            C_v_WC = W_v_WW = jnp.zeros(6)
+            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
+            C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
-            L_H_C = L_H_W = jax.vmap(
+            L_H_C = L_H_W = jax.vmap(  # noqa: F841
                 lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
             )(W_H_L)
 
-            L_v_LC = L_v_LW = jax.vmap(
+            L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
                 )
@@ -314,9 +314,9 @@ def bias_acceleration(
 
         case VelRepr.Mixed:
 
-            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL
+            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL  # noqa: F841
 
-            C_v_WC = LW_v_W_LW = jax.vmap(
+            C_v_WC = LW_v_W_LW = jax.vmap(  # noqa: F841
                 lambda i: js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Mixed
                 )
@@ -324,13 +324,13 @@ def bias_acceleration(
                 .set(jnp.zeros(3))
             )(jnp.arange(model.number_of_links()))
 
-            L_H_C = L_H_LW = jax.vmap(
+            L_H_C = L_H_LW = jax.vmap(  # noqa: F841
                 lambda W_H_L: jaxsim.math.Transform.inverse(
                     W_H_L.at[0:3, 3].set(jnp.zeros(3))
                 )
             )(W_H_L)
 
-            L_v_LC = L_v_L_LW = jax.vmap(
+            L_v_LC = L_v_L_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
                 )

jaxsim/api/contact.py

@@ -8,6 +8,7 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Cross, Transform
 from jaxsim.rbda.contacts.soft import SoftContactsParams
 
 from .common import VelRepr
@@ -411,3 +412,170 @@ def jacobian(
             raise ValueError(output_vel_repr)
 
     return O_J_WC
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def jacobian_derivative(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    r"""
+    Compute the derivative of the free-floating jacobian of the contact points.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian derivative.
+
+    Returns:
+        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the contact points.
+
+    Note:
+        The input representation of the free-floating jacobian derivative is the active
+        velocity representation.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the index of the parent link and the position of the collidable point.
+    parent_link_idxs = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+    L_p_Ci = jnp.array(model.kin_dyn_parameters.contact_parameters.point)
+    contact_idxs = jnp.arange(L_p_Ci.shape[0])
+
+    # Get the transforms of all the parent links.
+    W_H_Li = js.model.forward_kinematics(model=model, data=data)
+
+    # =====================================================
+    # Compute quantities to adjust the input representation
+    # =====================================================
+
+    def compute_T(model: js.model.JaxSimModel, X: jtp.Matrix) -> jtp.Matrix:
+        In = jnp.eye(model.dofs())
+        T = jax.scipy.linalg.block_diag(X, In)
+        return T
+
+    def compute_Ṫ(model: js.model.JaxSimModel, Ẋ: jtp.Matrix) -> jtp.Matrix:
+        On = jnp.zeros(shape=(model.dofs(), model.dofs()))
+        Ṫ = jax.scipy.linalg.block_diag(Ẋ, On)
+        return Ṫ
+
+    # Compute the operator to change the representation of ν, and its
+    # time derivative.
+    match data.velocity_representation:
+        case VelRepr.Inertial:
+            W_H_W = jnp.eye(4)
+            W_X_W = Adjoint.from_transform(transform=W_H_W)
+            W_Ẋ_W = jnp.zeros((6, 6))
+
+            T = compute_T(model=model, X=W_X_W)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
+
+        case VelRepr.Body:
+            W_H_B = data.base_transform()
+            W_X_B = Adjoint.from_transform(transform=W_H_B)
+            B_v_WB = data.base_velocity()
+            B_vx_WB = Cross.vx(B_v_WB)
+            W_Ẋ_B = W_X_B @ B_vx_WB
+
+            T = compute_T(model=model, X=W_X_B)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
+
+        case VelRepr.Mixed:
+            W_H_B = data.base_transform()
+            W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
+            W_X_BW = Adjoint.from_transform(transform=W_H_BW)
+            BW_v_WB = data.base_velocity()
+            BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
+            BW_vx_W_BW = Cross.vx(BW_v_W_BW)
+            W_Ẋ_BW = W_X_BW @ BW_vx_W_BW
+
+            T = compute_T(model=model, X=W_X_BW)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_BW)
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # =====================================================
+    # Compute quantities to adjust the output representation
+    # =====================================================
+
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        # Compute the Jacobian of the parent link in inertial representation.
+        W_J_WL_W = js.model.generalized_free_floating_jacobian(
+            model=model,
+            data=data,
+            output_vel_repr=VelRepr.Inertial,
+        )
+        # Compute the Jacobian derivative of the parent link in inertial representation.
+        W_J̇_WL_W = js.model.generalized_free_floating_jacobian_derivative(
+            model=model,
+            data=data,
+            output_vel_repr=VelRepr.Inertial,
+        )
+
+    # Get the Jacobian of the collidable points in the mixed representation.
+    with data.switch_velocity_representation(VelRepr.Mixed):
+        CW_J_WC_BW = jacobian(
+            model=model,
+            data=data,
+            output_vel_repr=VelRepr.Mixed,
+        )
+
+    def compute_O_J̇_WC_I(
+        L_p_C: jtp.Vector,
+        contact_idx: jtp.Int,
+        CW_J_WC_BW: jtp.Matrix,
+        W_H_L: jtp.Matrix,
+    ) -> jtp.Matrix:
+
+        parent_link_idx = parent_link_idxs[contact_idx]
+
+        match output_vel_repr:
+            case VelRepr.Inertial:
+                O_X_W = W_X_W = Adjoint.from_transform(  # noqa: F841
+                    transform=jnp.eye(4)
+                )
+                O_Ẋ_W = W_Ẋ_W = jnp.zeros((6, 6))  # noqa: F841
+
+            case VelRepr.Body:
+                L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
+                W_H_C = W_H_L[parent_link_idx] @ L_H_C
+                O_X_W = C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
+                with data.switch_velocity_representation(VelRepr.Inertial):
+                    W_nu = data.generalized_velocity()
+                W_v_WC = W_J_WL_W[parent_link_idx] @ W_nu
+                W_vx_WC = Cross.vx(W_v_WC)
+                O_Ẋ_W = C_Ẋ_W = -C_X_W @ W_vx_WC  # noqa: F841
+
+            case VelRepr.Mixed:
+                L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
+                W_H_C = W_H_L[parent_link_idx] @ L_H_C
+                W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
+                CW_H_W = Transform.inverse(W_H_CW)
+                O_X_W = CW_X_W = Adjoint.from_transform(transform=CW_H_W)
+                with data.switch_velocity_representation(VelRepr.Mixed):
+                    CW_v_WC = CW_J_WC_BW @ data.generalized_velocity()
+                W_v_W_CW = jnp.zeros(6).at[0:3].set(CW_v_WC[0:3])
+                W_vx_W_CW = Cross.vx(W_v_W_CW)
+                O_Ẋ_W = CW_Ẋ_W = -CW_X_W @ W_vx_W_CW  # noqa: F841
+
+            case _:
+                raise ValueError(output_vel_repr)
+
+        O_J̇_WC_I = jnp.zeros(shape=(6, 6 + model.dofs()))
+        O_J̇_WC_I += O_Ẋ_W @ W_J_WL_W[parent_link_idx] @ T
+        O_J̇_WC_I += O_X_W @ W_J̇_WL_W[parent_link_idx] @ T
+        O_J̇_WC_I += O_X_W @ W_J_WL_W[parent_link_idx] @ Ṫ
+
+        return O_J̇_WC_I
+
+    O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, 0, None))(
+        L_p_Ci, contact_idxs, CW_J_WC_BW, W_H_Li
+    )
+
+    return O_J̇_WC

jaxsim/api/data.py

@@ -12,7 +12,6 @@ import jaxlie
 import jaxsim.api as js
 import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.math import Quaternion
 from jaxsim.rbda.contacts.soft import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
@@ -191,9 +190,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_H_B = jaxlie.SE3.from_rotation_and_translation(
             translation=base_position,
-            rotation=jaxlie.SO3.from_quaternion_xyzw(
-                base_quaternion[jnp.array([1, 2, 3, 0])]
-            ),
+            rotation=jaxlie.SO3(wxyz=base_quaternion),
         ).as_matrix()
 
         v_WB = JaxSimModelData.other_representation_to_inertial(
@@ -380,13 +377,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
         )
 
-        return (
-            W_Q_B
-            if not dcm
-            else jaxlie.SO3.from_quaternion_xyzw(
-                Quaternion.to_xyzw(wxyz=W_Q_B)
-            ).as_matrix()
-        ).astype(float)
+        return (W_Q_B if not dcm else jaxlie.SO3(wxyz=W_Q_B).as_matrix()).astype(float)
 
     @jax.jit
     def base_transform(self) -> jtp.Matrix:

jaxsim/api/frame.py

@@ -7,7 +7,7 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim import exceptions
-from jaxsim.math import Adjoint, Transform
+from jaxsim.math import Adjoint, Cross, Transform
 
 from .common import VelRepr
 
@@ -257,3 +257,158 @@ def jacobian(
             raise ValueError(output_vel_repr)
 
     return O_J_WL_I
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def jacobian_derivative(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    r"""
+    Compute the derivative of the free-floating jacobian of the frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian derivative.
+
+    Returns:
+        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the frame.
+
+    Note:
+        The input representation of the free-floating jacobian derivative is the active
+        velocity representation.
+    """
+
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the index of the parent link.
+    L = idx_of_parent_link(model=model, frame_index=frame_index)
+
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        # Compute the Jacobian of the parent link in inertial representation.
+        W_J_WL_W = js.link.jacobian(
+            model=model,
+            data=data,
+            link_index=L,
+            output_vel_repr=VelRepr.Inertial,
+        )
+
+        # Compute the Jacobian derivative of the parent link in inertial representation.
+        W_J̇_WL_W = js.link.jacobian_derivative(
+            model=model,
+            data=data,
+            link_index=L,
+            output_vel_repr=VelRepr.Inertial,
+        )
+
+    # =====================================================
+    # Compute quantities to adjust the input representation
+    # =====================================================
+
+    def compute_T(model: js.model.JaxSimModel, X: jtp.Matrix) -> jtp.Matrix:
+        In = jnp.eye(model.dofs())
+        T = jax.scipy.linalg.block_diag(X, In)
+        return T
+
+    def compute_Ṫ(model: js.model.JaxSimModel, Ẋ: jtp.Matrix) -> jtp.Matrix:
+        On = jnp.zeros(shape=(model.dofs(), model.dofs()))
+        Ṫ = jax.scipy.linalg.block_diag(Ẋ, On)
+        return Ṫ
+
+    # Compute the operator to change the representation of ν, and its
+    # time derivative.
+    match data.velocity_representation:
+        case VelRepr.Inertial:
+            W_H_W = jnp.eye(4)
+            W_X_W = Adjoint.from_transform(transform=W_H_W)
+            W_Ẋ_W = jnp.zeros((6, 6))
+
+            T = compute_T(model=model, X=W_X_W)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
+
+        case VelRepr.Body:
+            W_H_B = data.base_transform()
+            W_X_B = Adjoint.from_transform(transform=W_H_B)
+            B_v_WB = data.base_velocity()
+            B_vx_WB = Cross.vx(B_v_WB)
+            W_Ẋ_B = W_X_B @ B_vx_WB
+
+            T = compute_T(model=model, X=W_X_B)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
+
+        case VelRepr.Mixed:
+            W_H_B = data.base_transform()
+            W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
+            W_X_BW = Adjoint.from_transform(transform=W_H_BW)
+            BW_v_WB = data.base_velocity()
+            BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
+            BW_vx_W_BW = Cross.vx(BW_v_W_BW)
+            W_Ẋ_BW = W_X_BW @ BW_vx_W_BW
+
+            T = compute_T(model=model, X=W_X_BW)
+            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_BW)
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # =====================================================
+    # Compute quantities to adjust the output representation
+    # =====================================================
+
+    match output_vel_repr:
+        case VelRepr.Inertial:
+            O_X_W = W_X_W = Adjoint.from_transform(transform=jnp.eye(4))
+            O_Ẋ_W = W_Ẋ_W = jnp.zeros((6, 6))
+
+        case VelRepr.Body:
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            O_X_W = F_X_W = Adjoint.from_transform(transform=W_H_F, inverse=True)
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_nu = data.generalized_velocity()
+            W_v_WF = W_J_WL_W @ W_nu
+            W_vx_WF = Cross.vx(W_v_WF)
+            O_Ẋ_W = F_Ẋ_W = -F_X_W @ W_vx_WF  # noqa: F841
+
+        case VelRepr.Mixed:
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            W_H_FW = W_H_F.at[0:3, 0:3].set(jnp.eye(3))
+            FW_H_W = Transform.inverse(W_H_FW)
+            O_X_W = FW_X_W = Adjoint.from_transform(transform=FW_H_W)
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                FW_J_WF_FW = jacobian(
+                    model=model,
+                    data=data,
+                    frame_index=frame_index,
+                    output_vel_repr=VelRepr.Mixed,
+                )
+                FW_v_WF = FW_J_WF_FW @ data.generalized_velocity()
+            W_v_W_FW = jnp.zeros(6).at[0:3].set(FW_v_WF[0:3])
+            W_vx_W_FW = Cross.vx(W_v_W_FW)
+            O_Ẋ_W = FW_Ẋ_W = -FW_X_W @ W_vx_W_FW  # noqa: F841
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    O_J̇_WF_I = jnp.zeros(shape=(6, 6 + model.dofs()))
+    O_J̇_WF_I += O_Ẋ_W @ W_J_WL_W @ T
+    O_J̇_WF_I += O_X_W @ W_J̇_WL_W @ T
+    O_J̇_WF_I += O_X_W @ W_J_WL_W @ Ṫ
+
+    return O_J̇_WF_I

jaxsim/api/link.py

@@ -288,7 +288,7 @@ def jacobian(
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
             B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
-            B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(
+            B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(  # noqa: F841
                 B_X_W, jnp.eye(model.dofs())
             )
 
@@ -299,7 +299,7 @@ def jacobian(
             W_R_B = data.base_orientation(dcm=True)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
             B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
-            B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(
+            B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(  # noqa: F841
                 B_X_BW, jnp.eye(model.dofs())
             )
 
@@ -313,7 +313,7 @@ def jacobian(
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
             W_X_B = Adjoint.from_transform(transform=W_H_B)
-            O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I
+            O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I  # noqa: F841
 
         case VelRepr.Body:
             L_X_B = Adjoint.from_transform(transform=B_H_L, inverse=True)
@@ -505,7 +505,7 @@ def jacobian_derivative(
             with data.switch_velocity_representation(VelRepr.Body):
                 B_v_WB = data.base_velocity()
 
-            O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)
+            O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)  # noqa: F841
 
         case VelRepr.Body:
 
@@ -519,7 +519,9 @@ def jacobian_derivative(
                 B_v_WB = data.base_velocity()
                 L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index)
 
-            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(B_X_L @ L_v_WL - B_v_WB)
+            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                B_X_L @ L_v_WL - B_v_WB
+            )
 
         case VelRepr.Mixed:
 
@@ -544,8 +546,9 @@ def jacobian_derivative(
             LW_v_LW_L = LW_v_WL - LW_v_W_LW
             LW_v_B_LW = LW_v_WL - LW_X_B @ B_v_WB - LW_v_LW_L
 
-            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(B_X_LW @ LW_v_B_LW)
-
+            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                B_X_LW @ LW_v_B_LW
+            )
         case _:
             raise ValueError(output_vel_repr)