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

jaxsim/api/frame.py

@@ -0,0 +1,471 @@
+import functools
+from collections.abc import Sequence
+
+import jax
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim import exceptions
+from jaxsim.math import Adjoint, Cross, Transform
+
+from .common import VelRepr
+
+# =======================
+# Index-related functions
+# =======================
+
+
+@jax.jit
+@js.common.named_scope
+def idx_of_parent_link(
+    model: js.model.JaxSimModel, *, frame_index: jtp.IntLike
+) -> jtp.Int:
+    """
+    Get the index of the link to which the frame is rigidly attached.
+
+    Args:
+        model: The model to consider.
+        frame_index: The index of the frame.
+
+    Returns:
+        The index of the frame's parent link.
+    """
+
+    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,
+    )
+
+    return jnp.array(model.kin_dyn_parameters.frame_parameters.body)[
+        frame_index - model.number_of_links()
+    ]
+
+
+@functools.partial(jax.jit, static_argnames="frame_name")
+@js.common.named_scope
+def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
+    """
+    Convert the name of a frame to its index.
+
+    Args:
+        model: The model to consider.
+        frame_name: The name of the frame.
+
+    Returns:
+        The index of the frame.
+    """
+
+    if frame_name not in model.frame_names():
+        raise ValueError(f"Frame '{frame_name}' not found in the model.")
+
+    return (
+        jnp.array(
+            model.number_of_links()
+            + model.kin_dyn_parameters.frame_parameters.name.index(frame_name)
+        )
+        .astype(int)
+        .squeeze()
+    )
+
+
+def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
+    """
+    Convert the index of a frame to its name.
+
+    Args:
+        model: The model to consider.
+        frame_index: The index of the frame.
+
+    Returns:
+        The name of the frame.
+    """
+
+    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,
+    )
+
+    return model.kin_dyn_parameters.frame_parameters.name[
+        frame_index - model.number_of_links()
+    ]
+
+
+@functools.partial(jax.jit, static_argnames=["frame_names"])
+@js.common.named_scope
+def names_to_idxs(
+    model: js.model.JaxSimModel, *, frame_names: Sequence[str]
+) -> jax.Array:
+    """
+    Convert a sequence of frame names to their corresponding indices.
+
+    Args:
+        model: The model to consider.
+        frame_names: The names of the frames.
+
+    Returns:
+        The indices of the frames.
+    """
+
+    return jnp.array(
+        [name_to_idx(model=model, frame_name=name) for name in frame_names]
+    ).astype(int)
+
+
+def idxs_to_names(
+    model: js.model.JaxSimModel, *, frame_indices: Sequence[jtp.IntLike]
+) -> tuple[str, ...]:
+    """
+    Convert a sequence of frame indices to their corresponding names.
+
+    Args:
+        model: The model to consider.
+        frame_indices: The indices of the frames.
+
+    Returns:
+        The names of the frames.
+    """
+
+    return tuple(idx_to_name(model=model, frame_index=idx) for idx in frame_indices)
+
+
+# ==========
+# Frame APIs
+# ==========
+
+
+@jax.jit
+@js.common.named_scope
+def transform(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+) -> jtp.Matrix:
+    """
+    Compute the SE(3) transform from the world frame to the specified frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame for which the transform is requested.
+
+    Returns:
+        The 4x4 matrix representing the transform.
+    """
+
+    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,
+    )
+
+    # Compute the necessary transforms.
+    L = idx_of_parent_link(model=model, frame_index=frame_index)
+    W_H_L = js.link.transform(model=model, data=data, link_index=L)
+
+    # Get the static frame pose wrt the parent link.
+    L_H_F = model.kin_dyn_parameters.frame_parameters.transform[
+        frame_index - model.number_of_links()
+    ]
+
+    # Combine the transforms computing the frame pose.
+    return W_H_L @ L_H_F
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@js.common.named_scope
+def velocity(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Vector:
+    """
+    Compute the 6D velocity 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 frame velocity.
+
+    Returns:
+        The 6D velocity of the frame in the specified velocity representation.
+    """
+    n_l = model.number_of_links()
+    n_f = model.number_of_frames()
+
+    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 frame jacobian having I as input representation (taken from data)
+    # and O as output representation, specified by the user (or taken from data).
+    O_J_WF_I = jacobian(
+        model=model,
+        data=data,
+        frame_index=frame_index,
+        output_vel_repr=output_vel_repr,
+    )
+
+    # Get the generalized velocity in the input velocity representation.
+    I_ν = data.generalized_velocity()
+
+    # Compute the frame velocity in the output velocity representation.
+    return O_J_WF_I @ I_ν
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@js.common.named_scope
+def jacobian(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    r"""
+    Compute 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.
+
+    Returns:
+        The :math:`6 \times (6+n)` free-floating jacobian of the frame.
+
+    Note:
+        The input representation of the free-floating jacobian is the active
+        velocity representation.
+    """
+
+    n_l = model.number_of_links()
+    n_f = model.number_of_frames()
+
+    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)
+
+    # Compute the Jacobian of the parent link using body-fixed output representation.
+    L_J_WL = js.link.jacobian(
+        model=model, data=data, link_index=L, output_vel_repr=VelRepr.Body
+    )
+
+    # Adjust the output representation.
+    match output_vel_repr:
+        case VelRepr.Inertial:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_X_L = Adjoint.from_transform(transform=W_H_L)
+            W_J_WL = W_X_L @ L_J_WL
+            O_J_WL_I = W_J_WL
+
+        case VelRepr.Body:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            F_H_L = Transform.inverse(W_H_F) @ W_H_L
+            F_X_L = Adjoint.from_transform(transform=F_H_L)
+            F_J_WL = F_X_L @ L_J_WL
+            O_J_WL_I = F_J_WL
+
+        case VelRepr.Mixed:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            F_H_L = Transform.inverse(W_H_F) @ W_H_L
+            FW_H_F = W_H_F.at[0:3, 3].set(jnp.zeros(3))
+            FW_H_L = FW_H_F @ F_H_L
+            FW_X_L = Adjoint.from_transform(transform=FW_H_L)
+            FW_J_WL = FW_X_L @ L_J_WL
+            O_J_WL_I = FW_J_WL
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    return O_J_WL_I
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@js.common.named_scope
+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