jaxsim 0.2.1.dev70__py3-none-any.whl → 0.2.1.dev80__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev70'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev70')
+__version__ = version = '0.2.1.dev80'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev80')

jaxsim/api/com.py

@@ -238,3 +238,179 @@ def average_centroidal_velocity_jacobian(
     G_Mbb = locked_centroidal_spatial_inertia(model=model, data=data)
 
     return jnp.linalg.inv(G_Mbb) @ G_J
+
+
+@jax.jit
+def bias_acceleration(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the bias linear acceleration of the center of mass.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The bias linear acceleration of the center of mass in the active representation.
+
+    Note:
+        The bias acceleration is expressed in the mixed frame
+        :math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
+        active velocity representation is either inertial-fixed or mixed,
+        and :math:`[C] = [B]` if the active velocity representation is body-fixed.
+    """
+
+    # Compute the pose of all links with forward kinematics.
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+
+    # Compute the bias acceleration of all links by zeroing the generalized velocity
+    # in the active representation.
+    v̇_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
+
+    def other_representation_to_body(
+        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
+        C_v̇_WL expressed in a generic frame C to the body-fixed representation L_v̇_WL.
+        """
+
+        L_X_C = jaxsim.math.Adjoint.from_transform(transform=L_H_C)
+        C_X_L = jaxsim.math.Adjoint.inverse(L_X_C)
+
+        L_v̇_WL = L_X_C @ (C_v̇_WL + jaxsim.math.Cross.vx(C_X_L @ L_v_LC) @ C_v_WC)
+        return L_v̇_WL
+
+    # We need here to get the body-fixed bias acceleration of the links.
+    # Since it's computed in the active representation, we need to convert it to body.
+    match data.velocity_representation:
+
+        case VelRepr.Body:
+            L_a_bias_WL = v̇_bias_WL
+
+        case VelRepr.Inertial:
+
+            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL
+            C_v_WC = W_v_WW = jnp.zeros(6)
+
+            L_H_C = L_H_W = jax.vmap(
+                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
+            )(W_H_L)
+
+            L_v_LC = L_v_LW = jax.vmap(
+                lambda i: -js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
+                )
+            )(jnp.arange(model.number_of_links()))
+
+            L_a_bias_WL = jax.vmap(
+                lambda i: other_representation_to_body(
+                    C_v̇_WL=C_v̇_WL[i],
+                    C_v_WC=C_v_WC,
+                    L_H_C=L_H_C[i],
+                    L_v_LC=L_v_LC[i],
+                )
+            )(jnp.arange(model.number_of_links()))
+
+        case VelRepr.Mixed:
+
+            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL
+
+            C_v_WC = LW_v_W_LW = jax.vmap(
+                lambda i: js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Mixed
+                )
+                .at[3:6]
+                .set(jnp.zeros(3))
+            )(jnp.arange(model.number_of_links()))
+
+            L_H_C = L_H_LW = jax.vmap(
+                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(
+                lambda i: -js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
+                )
+                .at[0:3]
+                .set(jnp.zeros(3))
+            )(jnp.arange(model.number_of_links()))
+
+            L_a_bias_WL = jax.vmap(
+                lambda i: other_representation_to_body(
+                    C_v̇_WL=C_v̇_WL[i],
+                    C_v_WC=C_v_WC[i],
+                    L_H_C=L_H_C[i],
+                    L_v_LC=L_v_LC[i],
+                )
+            )(jnp.arange(model.number_of_links()))
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Compute the bias of the 6D momentum derivative.
+    def bias_momentum_derivative_term(
+        link_index: jtp.Int, L_a_bias_WL: jtp.Vector
+    ) -> jtp.Vector:
+
+        # Get the body-fixed 6D inertia matrix.
+        L_M_L = js.link.spatial_inertia(model=model, link_index=link_index)
+
+        # Compute the body-fixed 6D velocity.
+        L_v_WL = js.link.velocity(
+            model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body
+        )
+
+        # Compute the world-to-link transformations for 6D forces.
+        W_Xf_L = jaxsim.math.Adjoint.from_transform(
+            transform=W_H_L[link_index], inverse=True
+        ).T
+
+        # Compute the contribution of the link to the bias acceleration of the CoM.
+        W_ḣ_bias_link_contribution = W_Xf_L @ (
+            L_M_L @ L_a_bias_WL + jaxsim.math.Cross.vx_star(L_v_WL) @ L_M_L @ L_v_WL
+        )
+
+        return W_ḣ_bias_link_contribution
+
+    # Sum the contributions of all links to the bias acceleration of the CoM.
+    W_ḣ_bias = jax.vmap(bias_momentum_derivative_term)(
+        jnp.arange(model.number_of_links()), L_a_bias_WL
+    ).sum(axis=0)
+
+    # Compute the total mass of the model.
+    m = js.model.total_mass(model=model)
+
+    # Compute the position of the CoM.
+    W_p_CoM = com_position(model=model, data=data)
+
+    match data.velocity_representation:
+
+        # G := G[W] = (W_p_CoM, [W])
+        case VelRepr.Inertial | VelRepr.Mixed:
+
+            W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+            GW_Xf_W = jaxsim.math.Adjoint.from_transform(W_H_GW).T
+
+            GW_ḣ_bias = GW_Xf_W @ W_ḣ_bias
+            GW_v̇l_com_bias = GW_ḣ_bias[0:3] / m
+
+            return GW_v̇l_com_bias
+
+        # G := G[B] = (W_p_CoM, [B])
+        case VelRepr.Body:
+
+            GB_Xf_W = jaxsim.math.Adjoint.from_transform(
+                transform=data.base_transform().at[0:3].set(W_p_CoM)
+            ).T
+
+            GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias
+            GB_v̇l_com_bias = GB_ḣ_bias[0:3] / m
+
+            return GB_v̇l_com_bias
+
+        case _:
+            raise ValueError(data.velocity_representation)

jaxsim/api/contact.py

@@ -32,6 +32,7 @@ def collidable_point_kinematics(
 
     from jaxsim.rbda import collidable_points
 
+    # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
     with data.switch_velocity_representation(VelRepr.Inertial):
         W_p_Ci, W_ṗ_Ci = collidable_points.collidable_points_pos_vel(
             model=model,
@@ -61,7 +62,9 @@ def collidable_point_positions(
         The position of the collidable points in the world frame.
     """
 
-    return collidable_point_kinematics(model=model, data=data)[0]
+    W_p_Ci, _ = collidable_point_kinematics(model=model, data=data)
+
+    return W_p_Ci
 
 
 @jax.jit
@@ -79,7 +82,9 @@ def collidable_point_velocities(
         The 3D velocity of the collidable points.
     """
 
-    return collidable_point_kinematics(model=model, data=data)[1]
+    _, W_ṗ_Ci = collidable_point_kinematics(model=model, data=data)
+
+    return W_ṗ_Ci
 
 
 @jax.jit
@@ -269,3 +274,124 @@ def estimate_good_soft_contacts_parameters(
     )
 
     return sc_parameters
+
+
+@jax.jit
+def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
+    r"""
+    Return the pose of the collidable points.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The stacked SE(3) matrices of all collidable points.
+
+    Note:
+        Each collidable point is implicitly associated with a frame
+        :math:`C = ({}^W p_C, [L])`, where :math:`{}^W p_C` is the position of the
+        collidable point and :math:`[L]` is the orientation frame of the link it is
+        rigidly attached to.
+    """
+
+    # Get the transforms of the parent link of all collidable points.
+    W_H_L = jax.vmap(
+        lambda parent_link_idx: js.link.transform(
+            model=model, data=data, link_index=parent_link_idx
+        )
+    )(jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int))
+
+    # Build the link-to-point transform from the displacement between the link frame L
+    # and the implicit contact frame C.
+    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(
+        model.kin_dyn_parameters.contact_parameters.point
+    )
+
+    # Compose the work-to-link and link-to-point transforms.
+    return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def jacobian(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Array:
+    r"""
+    Return the free-floating Jacobian of the collidable 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.
+
+    Returns:
+        The stacked 6×(6+n) free-floating jacobians of the frames associated to the
+        collidable points.
+
+    Note:
+        Each collidable point is implicitly associated with a frame
+        :math:`C = ({}^W p_C, [L])`, where :math:`{}^W p_C` is the position of the
+        collidable point and :math:`[L]` is the orientation frame of the link it is
+        rigidly attached to.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # For each collidable point, get the Jacobians of their parent link.
+    # In inertial-fixed output representation, the Jacobian of the parent link is also
+    # the Jacobian of the frame C implicitly associated with the collidable point.
+    W_J_WC = W_J_WL = jax.vmap(
+        lambda parent_link_idx: js.link.jacobian(
+            model=model,
+            data=data,
+            link_index=parent_link_idx,
+            output_vel_repr=VelRepr.Inertial,
+        )
+    )(jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int))
+
+    # Adjust the output representation.
+    match output_vel_repr:
+
+        case VelRepr.Inertial:
+            O_J_WC = W_J_WC
+
+        case VelRepr.Body:
+
+            W_H_C = transforms(model=model, data=data)
+
+            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+                C_X_W = jaxsim.math.Adjoint.from_transform(
+                    transform=W_H_C, inverse=True
+                )
+                C_J_WC = C_X_W @ W_J_WC
+                return C_J_WC
+
+            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+
+        case VelRepr.Mixed:
+
+            W_H_C = transforms(model=model, data=data)
+
+            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+
+                W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
+
+                CW_X_W = jaxsim.math.Adjoint.from_transform(
+                    transform=W_H_CW, inverse=True
+                )
+
+                CW_J_WC = CW_X_W @ W_J_WC
+                return CW_J_WC
+
+            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    return O_J_WC

{jaxsim-0.2.1.dev70.dist-info → jaxsim-0.2.1.dev80.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev70
+Version: 0.2.1.dev80
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev70.dist-info → jaxsim-0.2.1.dev80.dist-info}/RECORD

@@ -1,11 +1,11 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=1SWO53mcoalXzzBdA6_IrZRpZQsi_WIK6NLf67IAp-A,426
+jaxsim/_version.py,sha256=6bZIeO5gVEXbcUw14r0OEvahs0AZBEq5fAmJ16fr2fU,426
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
-jaxsim/api/com.py,sha256=Qtm_6qpiK4WtDVn6JMUHa8DySgBl9CjgKCybJqZ58Lc,7379
+jaxsim/api/com.py,sha256=Yof6otFi-mLWAs1rqjmeNJTOWIH9gn7BdU5EIjiL6Ts,13481
 jaxsim/api/common.py,sha256=DV-WZG28sikXopNv458aYvpLjmiAtFr5LRscOwXusuk,6640
-jaxsim/api/contact.py,sha256=Ve4ZOWkLEBRgK3KhtICxKY7YzsxYvc3lO-pPRBjqSnY,8659
+jaxsim/api/contact.py,sha256=Cvr-EfQtHP3nymtWdo-9WWU24Bkta-2Pp3nKsdjo6uc,12778
 jaxsim/api/data.py,sha256=GXSCfq4_PsRZuUwsmcl18e2Ppdtu4ykYs0oWih8J2ZI,26775
 jaxsim/api/frame.py,sha256=68liCmwFDikuP0GQOpo-nF7aVeHJ173e1USSOwskTP0,6246
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
@@ -58,8 +58,8 @@ jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
 jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim-0.2.1.dev70.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev70.dist-info/METADATA,sha256=4AdUyv7PhRO09XpEBc_79fG0MC-msgEQ-CnrstJesqA,9744
-jaxsim-0.2.1.dev70.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev70.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev70.dist-info/RECORD,,
+jaxsim-0.2.1.dev80.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev80.dist-info/METADATA,sha256=yFYgYOpZwSa-j1g4zj-4mp2B3z4TkL_9kbF-Rd5Q3hI,9744
+jaxsim-0.2.1.dev80.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev80.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev80.dist-info/RECORD,,