jaxsim 0.2.dev188__py3-none-any.whl → 0.2.dev364__py3-none-any.whl

jaxsim/api/link.py

@@ -4,20 +4,21 @@ from typing import Sequence
 import jax
 import jax.numpy as jnp
 import jaxlie
+import numpy as np
 
-import jaxsim.physics.algos.jacobian
+import jaxsim.api as js
+import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.high_level.common import VelRepr
 
-from . import data as Data
-from . import model as Model
+from .common import VelRepr
 
 # =======================
 # Index-related functions
 # =======================
 
 
-def name_to_idx(model: Model.JaxSimModel, *, link_name: str) -> jtp.Int:
+@functools.partial(jax.jit, static_argnames="link_name")
+def name_to_idx(model: js.model.JaxSimModel, *, link_name: str) -> jtp.Int:
     """
     Convert the name of a link to its index.
 
@@ -29,12 +30,18 @@ def name_to_idx(model: Model.JaxSimModel, *, link_name: str) -> jtp.Int:
         The index of the link.
     """
 
-    return jnp.array(
-        model.physics_model.description.links_dict[link_name].index, dtype=int
-    )
+    if link_name in model.kin_dyn_parameters.link_names:
+        return (
+            jnp.array(
+                np.argwhere(np.array(model.kin_dyn_parameters.link_names) == link_name)
+            )
+            .squeeze()
+            .astype(int)
+        )
+    return jnp.array(-1).astype(int)
 
 
-def idx_to_name(model: Model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
+def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
     """
     Convert the index of a link to its name.
 
@@ -46,11 +53,13 @@ def idx_to_name(model: Model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
         The name of the link.
     """
 
-    d = {l.index: l.name for l in model.physics_model.description.links_dict.values()}
-    return d[link_index]
+    return model.kin_dyn_parameters.link_names[link_index]
 
 
-def names_to_idxs(model: Model.JaxSimModel, *, link_names: Sequence[str]) -> jax.Array:
+@functools.partial(jax.jit, static_argnames="link_names")
+def names_to_idxs(
+    model: js.model.JaxSimModel, *, link_names: Sequence[str]
+) -> jax.Array:
     """
     Convert a sequence of link names to their corresponding indices.
 
@@ -63,13 +72,12 @@ def names_to_idxs(model: Model.JaxSimModel, *, link_names: Sequence[str]) -> jax
     """
 
     return jnp.array(
-        [model.physics_model.description.links_dict[name].index for name in link_names],
-        dtype=int,
-    )
+        [name_to_idx(model=model, link_name=name) for name in link_names],
+    ).astype(int)
 
 
 def idxs_to_names(
-    model: Model.JaxSimModel, *, link_indices: Sequence[jtp.IntLike] | jtp.VectorLike
+    model: js.model.JaxSimModel, *, link_indices: Sequence[jtp.IntLike] | jtp.VectorLike
 ) -> tuple[str, ...]:
     """
     Convert a sequence of link indices to their corresponding names.
@@ -82,8 +90,7 @@ def idxs_to_names(
         The names of the links.
     """
 
-    d = {l.index: l.name for l in model.physics_model.description.links_dict.values()}
-    return tuple(d[i] for i in link_indices)
+    return tuple(idx_to_name(model=model, link_index=idx) for idx in link_indices)
 
 
 # =========
@@ -91,21 +98,51 @@ def idxs_to_names(
 # =========
 
 
-def mass(model: Model.JaxSimModel, *, link_index: jtp.IntLike) -> jtp.Float:
-    """"""
+@jax.jit
+def mass(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> jtp.Float:
+    """
+    Return the mass of the link.
+
+    Args:
+        model: The model to consider.
+        link_index: The index of the link.
 
-    return model.physics_model._link_masses[link_index].astype(float)
+    Returns:
+        The mass of the link.
+    """
 
+    return model.kin_dyn_parameters.link_parameters.mass[link_index].astype(float)
 
-def spatial_inertia(model: Model.JaxSimModel, *, link_index: jtp.IntLike) -> jtp.Matrix:
-    """"""
 
-    return model.physics_model._link_spatial_inertias[link_index]
+@jax.jit
+def spatial_inertia(
+    model: js.model.JaxSimModel, *, link_index: jtp.IntLike
+) -> jtp.Matrix:
+    """
+    Compute the 6D spatial inertial of the link.
+
+    Args:
+        model: The model to consider.
+        link_index: The index of the link.
+
+    Returns:
+        The 6×6 matrix representing the spatial inertia of the link expressed in
+        the link frame (body-fixed representation).
+    """
+
+    link_parameters = jax.tree_util.tree_map(
+        lambda l: l[link_index], model.kin_dyn_parameters.link_parameters
+    )
+
+    return js.kin_dyn_parameters.LinkParameters.spatial_inertia(link_parameters)
 
 
 @jax.jit
 def transform(
-    model: Model.JaxSimModel, data: Data.JaxSimModelData, *, link_index: jtp.IntLike
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_index: jtp.IntLike,
 ) -> jtp.Matrix:
     """
     Compute the SE(3) transform from the world frame to the link frame.
@@ -119,13 +156,13 @@ def transform(
         The 4x4 matrix representing the transform.
     """
 
-    return Model.forward_kinematics(model=model, data=data)[link_index]
+    return js.model.forward_kinematics(model=model, data=data)[link_index]
 
 
 @jax.jit
 def com_position(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
     *,
     link_index: jtp.IntLike,
     in_link_frame: jtp.BoolLike = True,
@@ -168,8 +205,8 @@ def com_position(
 
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
     *,
     link_index: jtp.IntLike,
     output_vel_repr: VelRepr | None = None,
@@ -185,78 +222,116 @@ def jacobian(
             The output velocity representation of the free-floating jacobian.
 
     Returns:
-        The 6x(6+dofs) free-floating jacobian of the link.
+        The 6×(6+n) free-floating jacobian of the link.
 
     Note:
         The input representation of the free-floating jacobian is the active
         velocity representation.
     """
 
-    if output_vel_repr is None:
-        output_vel_repr = data.velocity_representation
-
-    # Compute the doubly left-trivialized free-floating jacobian
-    L_J_WL_B = jaxsim.physics.algos.jacobian.jacobian(
-        model=model.physics_model,
-        body_index=link_index,
-        q=data.joint_positions(),
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
-    match data.velocity_representation:
+    # Compute the doubly-left free-floating full jacobian.
+    B_J_full_WX_B, B_H_Li = jaxsim.rbda.jacobian_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+    )
 
-        case VelRepr.Body:
-            L_J_WL_target = L_J_WL_B
+    # Compute the actual doubly-left free-floating jacobian of the link.
+    κ = model.kin_dyn_parameters.support_body_array_bool[link_index]
+    B_J_WL_B = jnp.hstack([jnp.ones(5), κ]) * B_J_full_WX_B
 
+    # Adjust the input representation such that `J_WL_I @ I_ν`.
+    match data.velocity_representation:
         case VelRepr.Inertial:
-            dofs = model.dofs()
             W_H_B = data.base_transform()
-
             B_X_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
-            zero_6n = jnp.zeros(shape=(6, dofs))
-
-            B_T_W = jnp.vstack(
-                [
-                    jnp.block([B_X_W, zero_6n]),
-                    jnp.block([zero_6n.T, jnp.eye(dofs)]),
-                ]
+            B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(
+                B_X_W, jnp.eye(model.dofs())
             )
 
-            L_J_WL_target = L_J_WL_B @ B_T_W
+        case VelRepr.Body:
+            B_J_WL_I = B_J_WL_B
 
         case VelRepr.Mixed:
-            dofs = model.dofs()
-            W_H_B = data.base_transform()
-            BW_H_B = jnp.array(W_H_B).at[0:3, 3].set(jnp.zeros(3))
-
+            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 = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
-            zero_6n = jnp.zeros(shape=(6, dofs))
-
-            B_T_BW = jnp.vstack(
-                [
-                    jnp.block([B_X_BW, zero_6n]),
-                    jnp.block([zero_6n.T, jnp.eye(dofs)]),
-                ]
+            B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(
+                B_X_BW, jnp.eye(model.dofs())
             )
 
-            L_J_WL_target = L_J_WL_B @ B_T_BW
-
         case _:
             raise ValueError(data.velocity_representation)
 
-    match output_vel_repr:
-        case VelRepr.Body:
-            return L_J_WL_target
+    B_H_L = B_H_Li[link_index]
 
+    # Adjust the output representation such that `O_v_WL_I = O_J_WL_I @ I_ν`.
+    match output_vel_repr:
         case VelRepr.Inertial:
-            W_H_L = transform(model=model, data=data, link_index=link_index)
-            W_X_L = jaxlie.SE3.from_matrix(W_H_L).adjoint()
-            return W_X_L @ L_J_WL_target
+            W_H_B = data.base_transform()
+            W_X_B = jaxlie.SE3.from_matrix(W_H_B).adjoint()
+            O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I
+
+        case VelRepr.Body:
+            L_X_B = jaxlie.SE3.from_matrix(B_H_L).inverse().adjoint()
+            L_J_WL_I = L_X_B @ B_J_WL_I
+            O_J_WL_I = L_J_WL_I
 
         case VelRepr.Mixed:
-            W_H_L = transform(model=model, data=data, link_index=link_index)
-            LW_H_L = jnp.array(W_H_L).at[0:3, 3].set(jnp.zeros(3))
-            LW_X_L = jaxlie.SE3.from_matrix(LW_H_L).adjoint()
-            return LW_X_L @ L_J_WL_target
+            W_H_B = data.base_transform()
+            W_H_L = W_H_B @ B_H_L
+            LW_H_L = W_H_L.at[0:3, 3].set(jnp.zeros(3))
+            LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
+            LW_X_B = jaxlie.SE3.from_matrix(LW_H_B).adjoint()
+            LW_J_WL_I = LW_X_B @ B_J_WL_I
+            O_J_WL_I = LW_J_WL_I
 
         case _:
             raise ValueError(output_vel_repr)
+
+    return O_J_WL_I
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def velocity(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Vector:
+    """
+    Compute the 6D velocity of the link.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_index: The index of the link.
+        output_vel_repr:
+            The output velocity representation of the link velocity.
+
+    Returns:
+        The 6D velocity of the link in the specified velocity representation.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the link jacobian having I as input representation (taken from data)
+    # and O as output representation, specified by the user (or taken from data).
+    O_J_WL_I = jacobian(
+        model=model,
+        data=data,
+        link_index=link_index,
+        output_vel_repr=output_vel_repr,
+    )
+
+    # Get the generalized velocity in the input velocity representation.
+    I_ν = data.generalized_velocity()
+
+    # Compute the link velocity in the output velocity representation.
+    return O_J_WL_I @ I_ν