jaxsim 0.6.2.dev2__py3-none-any.whl → 0.6.2.dev102__py3-none-any.whl

jaxsim/api/frame.py

@@ -229,7 +229,7 @@ def velocity(
     )
 
     # Get the generalized velocity in the input velocity representation.
-    I_ν = data.generalized_velocity()
+    I_ν = data.generalized_velocity
 
     # Compute the frame velocity in the output velocity representation.
     return O_J_WF_I @ I_ν
@@ -401,9 +401,9 @@ def jacobian_derivative(
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
 
         case VelRepr.Body:
-            W_H_B = data.base_transform()
+            W_H_B = data._base_transform
             W_X_B = Adjoint.from_transform(transform=W_H_B)
-            B_v_WB = data.base_velocity()
+            B_v_WB = data.base_velocity
             B_vx_WB = Cross.vx(B_v_WB)
             W_Ẋ_B = W_X_B @ B_vx_WB
 
@@ -411,10 +411,10 @@ def jacobian_derivative(
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
 
         case VelRepr.Mixed:
-            W_H_B = data.base_transform()
+            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_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
@@ -438,7 +438,7 @@ def jacobian_derivative(
             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_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
@@ -455,7 +455,7 @@ def jacobian_derivative(
                     frame_index=frame_index,
                     output_vel_repr=VelRepr.Mixed,
                 )
-                FW_v_WF = FW_J_WF_FW @ data.generalized_velocity()
+                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

jaxsim/api/integrators.py

@@ -0,0 +1,76 @@
+import dataclasses
+
+import jax.numpy as jnp
+
+import jaxsim
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.api.data import JaxSimModelData
+from jaxsim.math import Adjoint, Transform
+
+
+def semi_implicit_euler_integration(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    base_acceleration_inertial: jtp.Vector,
+    joint_accelerations: jtp.Vector,
+) -> JaxSimModelData:
+    """Integrate the system state using the semi-implicit Euler method."""
+    # Step the dynamics forward.
+    with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
+
+        dt = model.time_step
+        W_v̇_WB = base_acceleration_inertial
+        s̈ = joint_accelerations
+
+        B_H_W = Transform.inverse(data._base_transform).at[:3, :3].set(jnp.eye(3))
+        BW_X_W = Adjoint.from_transform(B_H_W)
+
+        new_generalized_acceleration = jnp.hstack([W_v̇_WB, s̈])
+
+        new_generalized_velocity = (
+            data.generalized_velocity + dt * new_generalized_acceleration
+        )
+
+        new_base_velocity_inertial = new_generalized_velocity[0:6]
+        new_joint_velocities = new_generalized_velocity[6:]
+
+        base_lin_velocity_inertial = new_base_velocity_inertial[0:3]
+
+        new_base_velocity_mixed = BW_X_W @ new_generalized_velocity[0:6]
+        base_lin_velocity_mixed = new_base_velocity_mixed[0:3]
+        base_ang_velocity_mixed = new_base_velocity_mixed[3:6]
+
+        base_quaternion_derivative = jaxsim.math.Quaternion.derivative(
+            quaternion=data.base_orientation,
+            omega=base_ang_velocity_mixed,
+            omega_in_body_fixed=False,
+        ).squeeze()
+
+        new_base_position = data.base_position + dt * base_lin_velocity_mixed
+        new_base_quaternion = data.base_orientation + dt * base_quaternion_derivative
+
+        base_quaternion_norm = jaxsim.math.safe_norm(new_base_quaternion)
+
+        new_base_quaternion = new_base_quaternion / jnp.where(
+            base_quaternion_norm == 0, 1.0, base_quaternion_norm
+        )
+
+        new_joint_position = data.joint_positions + dt * new_joint_velocities
+
+    # TODO: Avoid double replace, e.g. by computing cached value here
+    data = dataclasses.replace(
+        data,
+        _base_quaternion=new_base_quaternion,
+        _base_position=new_base_position,
+        _joint_positions=new_joint_position,
+        _joint_velocities=new_joint_velocities,
+        _base_linear_velocity=base_lin_velocity_inertial,
+        # Here we use the base angular velocity in mixed representation since
+        # it's equivalent to the one in inertial representation
+        # See: S. Traversaro and A. Saccon, “Multibody Dynamics Notation (Version 2), pg.9
+        _base_angular_velocity=base_ang_velocity_mixed,
+    )
+    data = data.replace(model=model)  # update cache
+
+    return data

jaxsim/api/kin_dyn_parameters.py

@@ -11,7 +11,7 @@ from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
 from jaxsim.math import Adjoint, Inertia, JointModel, supported_joint_motion
-from jaxsim.parsers.descriptions import JointDescription, ModelDescription
+from jaxsim.parsers.descriptions import JointDescription, JointType, ModelDescription
 from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
 
@@ -36,6 +36,7 @@ class KinDynParameters(JaxsimDataclass):
     link_names: Static[tuple[str]]
     _parent_array: Static[HashedNumpyArray]
     _support_body_array_bool: Static[HashedNumpyArray]
+    _motion_subspaces: Static[HashedNumpyArray]
 
     # Links
     link_parameters: LinkParameters
@@ -50,6 +51,13 @@ class KinDynParameters(JaxsimDataclass):
     joint_model: JointModel
     joint_parameters: JointParameters | None
 
+    @property
+    def motion_subspaces(self) -> jtp.Matrix:
+        r"""
+        Return the motion subspaces :math:`\mathbf{S}(s)` of the joints.
+        """
+        return self._motion_subspaces.get()
+
     @property
     def parent_array(self) -> jtp.Vector:
         r"""
@@ -215,6 +223,31 @@ class KinDynParameters(JaxsimDataclass):
             jnp.arange(start=0, stop=len(ordered_links))
         )
 
+        def motion_subspace(joint_type: int, axis: npt.ArrayLike) -> npt.ArrayLike:
+
+            S = {
+                JointType.Fixed: np.zeros(shape=(6, 1)),
+                JointType.Revolute: np.vstack(np.hstack([np.zeros(3), axis.axis])),
+                JointType.Prismatic: np.vstack(np.hstack([axis.axis, np.zeros(3)])),
+            }
+
+            return S[joint_type]
+
+        S_J = (
+            jnp.array(
+                [
+                    motion_subspace(joint_type, axis)
+                    for joint_type, axis in zip(
+                        joint_model.joint_types[1:], joint_model.joint_axis, strict=True
+                    )
+                ]
+            )
+            if len(joint_model.joint_axis) != 0
+            else jnp.empty((0, 6, 1))
+        )
+
+        motion_subspaces = jnp.vstack([jnp.zeros((6, 1))[jnp.newaxis, ...], S_J])
+
         # =================================
         # Build and return KinDynParameters
         # =================================
@@ -223,6 +256,7 @@ class KinDynParameters(JaxsimDataclass):
             link_names=tuple(l.name for l in ordered_links),
             _parent_array=HashedNumpyArray(array=parent_array),
             _support_body_array_bool=HashedNumpyArray(array=support_body_array_bool),
+            _motion_subspaces=HashedNumpyArray(array=motion_subspaces),
             link_parameters=link_parameters,
             joint_model=joint_model,
             joint_parameters=joint_parameters,
@@ -359,54 +393,6 @@ class KinDynParameters(JaxsimDataclass):
             of each joint.
         """
 
-        return self.joint_transforms_and_motion_subspaces(
-            joint_positions=joint_positions,
-            base_transform=base_transform,
-        )[0]
-
-    @jax.jit
-    def joint_motion_subspaces(
-        self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
-    ) -> jtp.Array:
-        r"""
-        Return the motion subspaces of the joints.
-
-        Args:
-            joint_positions: The joint positions.
-            base_transform: The homogeneous matrix defining the base pose.
-
-        Returns:
-            The stacked motion subspaces :math:`\mathbf{S}(s)` of each joint.
-        """
-
-        return self.joint_transforms_and_motion_subspaces(
-            joint_positions=joint_positions,
-            base_transform=base_transform,
-        )[1]
-
-    @jax.jit
-    def joint_transforms_and_motion_subspaces(
-        self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
-    ) -> tuple[jtp.Array, jtp.Array]:
-        r"""
-        Return the transforms and the motion subspaces of the joints.
-
-        Args:
-            joint_positions: The joint positions.
-            base_transform: The homogeneous matrix defining the base pose.
-
-        Returns:
-            A tuple containing the stacked transforms
-            :math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
-            and the stacked motion subspaces :math:`\mathbf{S}(s)` of each joint.
-
-        Note:
-            The first transform, at index 0, provides the pose of the base link
-            w.r.t. the world frame. For both floating-base and fixed-base systems,
-            it takes into account the base pose and the optional transform
-            between the root frame of the model and the base link.
-        """
-
         # Rename the base transform.
         W_H_B = base_transform
 
@@ -417,22 +403,19 @@ class KinDynParameters(JaxsimDataclass):
                 self.joint_model.λ_H_pre[1 : 1 + self.number_of_joints()],
             ]
         )
-
-        # Compute the transforms and motion subspaces of the joints.
         if self.number_of_joints() == 0:
-            pre_H_suc_J, S_J = jnp.empty((0, 4, 4)), jnp.empty((0, 6, 1))
+            pre_H_suc_J = jnp.empty((0, 4, 4))
         else:
-            pre_H_suc_J, S_J = jax.vmap(supported_joint_motion)(
-                jnp.array(self.joint_model.joint_types[1:]).astype(int),
-                jnp.array(joint_positions),
-                jnp.array([j.axis for j in self.joint_model.joint_axis]),
+            pre_H_suc_J = jax.vmap(supported_joint_motion)(
+                joint_types=jnp.array(self.joint_model.joint_types[1:]).astype(int),
+                joint_positions=jnp.array(joint_positions),
+                joint_axes=jnp.array([j.axis for j in self.joint_model.joint_axis]),
             )
 
         # Extract the transforms and motion subspaces of the joints.
         # We stack the base transform W_H_B at index 0, and a dummy motion subspace
         # for either the fixed or free-floating joint connecting the world to the base.
         pre_H_suc = jnp.vstack([W_H_B[jnp.newaxis, ...], pre_H_suc_J])
-        S = jnp.vstack([jnp.zeros((6, 1))[jnp.newaxis, ...], S_J])
 
         # Extract the successor-to-child fixed transforms.
         # Note that here we include also the index 0 since suc_H_child[0] stores the
@@ -448,7 +431,7 @@ class KinDynParameters(JaxsimDataclass):
             )
         )(λ_H_pre, pre_H_suc, suc_H_i)
 
-        return i_X_λ, S
+        return i_X_λ
 
     # ============================
     # Helpers to update parameters

jaxsim/api/link.py

@@ -187,7 +187,7 @@ def transform(
         idx=link_index,
     )
 
-    return js.model.forward_kinematics(model=model, data=data)[link_index]
+    return data._link_transforms[link_index]
 
 
 @jax.jit
@@ -275,7 +275,7 @@ def jacobian(
     # 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(),
+        joint_positions=data.joint_positions,
     )
 
     # Compute the actual doubly-left free-floating jacobian of the link.
@@ -285,7 +285,7 @@ def jacobian(
     # Adjust the input representation such that `J_WL_I @ I_ν`.
     match data.velocity_representation:
         case VelRepr.Inertial:
-            W_H_B = data.base_transform()
+            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(  # noqa: F841
                 B_X_W, jnp.eye(model.dofs())
@@ -295,7 +295,7 @@ def jacobian(
             B_J_WL_I = B_J_WL_B
 
         case VelRepr.Mixed:
-            W_R_B = data.base_orientation(dcm=True)
+            W_R_B = jaxsim.math.Quaternion.to_dcm(data.base_orientation)
             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(  # noqa: F841
@@ -310,7 +310,7 @@ def jacobian(
     # Adjust the output representation such that `O_v_WL_I = O_J_WL_I @ I_ν`.
     match output_vel_repr:
         case VelRepr.Inertial:
-            W_H_B = data.base_transform()
+            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  # noqa: F841
 
@@ -320,7 +320,7 @@ def jacobian(
             O_J_WL_I = L_J_WL_I
 
         case VelRepr.Mixed:
-            W_H_B = data.base_transform()
+            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)
@@ -378,7 +378,7 @@ def velocity(
     )
 
     # Get the generalized velocity in the input velocity representation.
-    I_ν = data.generalized_velocity()
+    I_ν = data.generalized_velocity
 
     # Compute the link velocity in the output velocity representation.
     return O_J_WL_I @ I_ν