jaxsim 0.2.dev191__py3-none-any.whl → 0.2.dev366__py3-none-any.whl

jaxsim/__init__.py

@@ -61,7 +61,6 @@ del _jnp_options
 del _np_options
 del _is_editable
 
-from . import high_level, logging, math, simulation, sixd
-from .high_level.common import VelRepr
-from .simulation.ode_integration import IntegratorType
-from .simulation.simulator import JaxSim
+from . import terrain  # isort:skip
+from . import api, integrators, logging, math, rbda
+from .api.common import VelRepr

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.dev191'
-__version_tuple__ = version_tuple = (0, 2, 'dev191')
+__version__ = version = '0.2.dev366'
+__version_tuple__ = version_tuple = (0, 2, 'dev366')

jaxsim/api/__init__.py

@@ -1 +1,3 @@
-from . import contact, data, joint, link, model, ode
+from . import common  # isort:skip
+from . import model, data  # isort:skip
+from . import com, contact, joint, kin_dyn_parameters, link, ode, ode_data, references

jaxsim/api/com.py

@@ -0,0 +1,240 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.math
+import jaxsim.typing as jtp
+
+from .common import VelRepr
+
+
+@jax.jit
+def com_position(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    """
+    Compute the position of the center of mass of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The position of the center of mass of the model w.r.t. the world frame.
+    """
+
+    m = js.model.total_mass(model=model)
+
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+    W_H_B = data.base_transform()
+    B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
+
+    def B_p̃_LCoM(i) -> jtp.Vector:
+        m = js.link.mass(model=model, link_index=i)
+        L_p_LCoM = js.link.com_position(
+            model=model, data=data, link_index=i, in_link_frame=True
+        )
+        return m * B_H_W @ W_H_L[i] @ jnp.hstack([L_p_LCoM, 1])
+
+    com_links = jax.vmap(B_p̃_LCoM)(jnp.arange(model.number_of_links()))
+
+    B_p̃_CoM = (1 / m) * com_links.sum(axis=0)
+    B_p̃_CoM = B_p̃_CoM.at[3].set(1)
+
+    return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
+
+
+@jax.jit
+def com_linear_velocity(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the linear velocity of the center of mass of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The linear velocity of the center of mass of the model in the
+        active representation.
+
+    Note:
+        The linear velocity of the center of mass  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.
+    """
+
+    # Extract the linear component of the 6D average centroidal velocity.
+    # This is expressed in G[B] in body-fixed representation, and in G[W] in
+    # inertial-fixed or mixed representation.
+    G_vl_WG = average_centroidal_velocity(model=model, data=data)[0:3]
+
+    return G_vl_WG
+
+
+@jax.jit
+def centroidal_momentum(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the centroidal momentum of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The centroidal momentum of the model.
+
+    Note:
+        The centroidal momentum is expressed in the mixed frame
+        :math:`({}^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.
+    """
+
+    ν = data.generalized_velocity()
+    G_J = centroidal_momentum_jacobian(model=model, data=data)
+
+    return G_J @ ν
+
+
+@jax.jit
+def centroidal_momentum_jacobian(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    r"""
+    Compute the Jacobian of the centroidal momentum of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The Jacobian of the centroidal momentum of the model.
+
+    Note:
+        The frame corresponding to the output representation of this Jacobian is either
+        :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
+        or :math:`G[B]`, if the active velocity representation is body-fixed.
+
+    Note:
+        This Jacobian is also known in the literature as Centroidal Momentum Matrix.
+    """
+
+    # Compute the Jacobian of the total momentum with body-fixed output representation.
+    # We convert the output representation either to G[W] or G[B] below.
+    B_Jh = js.model.total_momentum_jacobian(
+        model=model, data=data, output_vel_repr=VelRepr.Body
+    )
+
+    W_H_B = data.base_transform()
+    B_H_W = jaxsim.math.Transform.inverse(W_H_B)
+
+    W_p_CoM = com_position(model=model, data=data)
+
+    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)
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Compute the transform for 6D forces.
+    G_Xf_B = jaxsim.math.Adjoint.from_transform(transform=B_H_W @ W_H_G).T
+
+    return G_Xf_B @ B_Jh
+
+
+@jax.jit
+def locked_centroidal_spatial_inertia(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+):
+    """
+    Compute the locked centroidal spatial inertia of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The locked centroidal spatial inertia of the model.
+    """
+
+    with data.switch_velocity_representation(VelRepr.Body):
+        B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
+
+    W_H_B = data.base_transform()
+    W_p_CoM = com_position(model=model, data=data)
+
+    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)
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    B_H_G = jaxlie.SE3.from_matrix(jaxsim.math.Transform.inverse(W_H_B) @ W_H_G)
+
+    B_Xv_G = B_H_G.adjoint()
+    G_Xf_B = B_Xv_G.transpose()
+
+    return G_Xf_B @ B_Mbb_B @ B_Xv_G
+
+
+@jax.jit
+def average_centroidal_velocity(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the average centroidal velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The average centroidal velocity of the model.
+
+    Note:
+        The average velocity 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.
+    """
+
+    ν = data.generalized_velocity()
+    G_J = average_centroidal_velocity_jacobian(model=model, data=data)
+
+    return G_J @ ν
+
+
+@jax.jit
+def average_centroidal_velocity_jacobian(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    r"""
+    Compute the Jacobian of the average centroidal velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The Jacobian of the average centroidal velocity of the model.
+
+    Note:
+        The frame corresponding to the output representation of this Jacobian is either
+        :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
+        or :math:`G[B]`, if the active velocity representation is body-fixed.
+    """
+
+    G_J = centroidal_momentum_jacobian(model=model, data=data)
+    G_Mbb = locked_centroidal_spatial_inertia(model=model, data=data)
+
+    return jnp.linalg.inv(G_Mbb) @ G_J

jaxsim/api/common.py

@@ -1,6 +1,7 @@
 import abc
 import contextlib
 import dataclasses
+import enum
 import functools
 from typing import ContextManager
 
@@ -11,7 +12,6 @@ import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
-from jaxsim.high_level.common import VelRepr
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 try:
@@ -20,6 +20,17 @@ except ImportError:
     from typing_extensions import Self
 
 
+@enum.unique
+class VelRepr(enum.IntEnum):
+    """
+    Enumeration of all supported 6D velocity representations.
+    """
+
+    Body = enum.auto()
+    Mixed = enum.auto()
+    Inertial = enum.auto()
+
+
 @jax_dataclasses.pytree_dataclass
 class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
     """
@@ -59,7 +70,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             # We run this in a mutable context with restoration so that any exception
             # occurring, we restore the original object in case it was modified.
             with self.mutable_context(
-                mutability=self._mutability(), restore_after_exception=True
+                mutability=self.mutability(), restore_after_exception=True
             ):
                 yield self
 

jaxsim/api/contact.py

@@ -3,16 +3,16 @@ import functools
 import jax
 import jax.numpy as jnp
 
+import jaxsim.api as js
+import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.physics.algos import soft_contacts
 
-from . import data as Data
-from . import model as Model
+from .common import VelRepr
 
 
 @jax.jit
 def collidable_point_kinematics(
-    model: Model.JaxSimModel, data: Data.JaxSimModelData
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
 ) -> tuple[jtp.Matrix, jtp.Matrix]:
     """
     Compute the position and 3D velocity of the collidable points in the world frame.
@@ -30,21 +30,25 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    from jaxsim.physics.algos.soft_contacts import collidable_points_pos_vel
-
-    W_p_Ci, W_ṗ_Ci = collidable_points_pos_vel(
-        model=model.physics_model,
-        q=data.state.physics_model.joint_positions,
-        qd=data.state.physics_model.joint_velocities,
-        xfb=data.state.physics_model.xfb(),
-    )
+    from jaxsim.rbda import collidable_points
+
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        W_p_Ci, W_ṗ_Ci = collidable_points.collidable_points_pos_vel(
+            model=model,
+            base_position=data.base_position(),
+            base_quaternion=data.base_orientation(dcm=False),
+            joint_positions=data.joint_positions(model=model),
+            base_linear_velocity=data.base_velocity()[0:3],
+            base_angular_velocity=data.base_velocity()[3:6],
+            joint_velocities=data.joint_velocities(model=model),
+        )
 
-    return W_p_Ci.T, W_ṗ_Ci.T
+    return W_p_Ci, W_ṗ_Ci
 
 
 @jax.jit
 def collidable_point_positions(
-    model: Model.JaxSimModel, data: Data.JaxSimModelData
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
 ) -> jtp.Matrix:
     """
     Compute the position of the collidable points in the world frame.
@@ -62,7 +66,7 @@ def collidable_point_positions(
 
 @jax.jit
 def collidable_point_velocities(
-    model: Model.JaxSimModel, data: Data.JaxSimModelData
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
 ) -> jtp.Matrix:
     """
     Compute the 3D velocity of the collidable points in the world frame.
@@ -78,10 +82,83 @@ def collidable_point_velocities(
     return collidable_point_kinematics(model=model, data=data)[1]
 
 
+@jax.jit
+def collidable_point_forces(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    """
+    Compute the 6D forces applied to each collidable point.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The 6D forces applied to each collidable point expressed in the frame
+        corresponding to the active representation.
+    """
+
+    f_Ci, _ = collidable_point_dynamics(model=model, data=data)
+
+    return f_Ci
+
+
+@jax.jit
+def collidable_point_dynamics(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> tuple[jtp.Matrix, jtp.Matrix]:
+    r"""
+    Compute the 6D force applied to each collidable point and the corresponding
+    material deformation rate.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The 6D force applied to each collidable point and the corresponding
+        material deformation rate.
+
+    Note:
+        The material deformation rate is always returned in the mixed frame
+        `C[W] = ({}^W \mathbf{p}_C, [W])`. This is convenient for integration purpose.
+        Instead, the 6D forces are returned in the active representation.
+    """
+
+    # Compute the position and linear velocities (mixed representation) of
+    # all collidable points belonging to the robot.
+    W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
+
+    # Build the soft contact model.
+    soft_contacts = jaxsim.rbda.SoftContacts(
+        parameters=data.soft_contacts_params, terrain=model.terrain
+    )
+
+    # Compute the 6D force expressed in the inertial frame and applied to each
+    # collidable point, and the corresponding material deformation rate.
+    # Note that the material deformation rate is always returned in the mixed frame
+    # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
+    W_f_Ci, CW_ṁ = jax.vmap(soft_contacts.contact_model)(
+        W_p_Ci, W_ṗ_Ci, data.state.soft_contacts.tangential_deformation
+    )
+
+    # Convert the 6D forces to the active representation.
+    f_Ci = jax.vmap(
+        lambda W_f_C: data.inertial_to_other_representation(
+            array=W_f_C,
+            other_representation=data.velocity_representation,
+            transform=data.base_transform(),
+            is_force=True,
+        )
+    )(W_f_Ci)
+
+    return f_Ci, CW_ṁ
+
+
 @functools.partial(jax.jit, static_argnames=["link_names"])
 def in_contact(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
     *,
     link_names: tuple[str, ...] | None = None,
 ) -> jtp.Vector:
@@ -100,48 +177,44 @@ def in_contact(
 
     link_names = link_names if link_names is not None else model.link_names()
 
-    if set(link_names) - set(model.link_names()) != set():
+    if set(link_names).difference(model.link_names()):
         raise ValueError("One or more link names are not part of the model")
 
-    from jaxsim.physics.algos.soft_contacts import collidable_points_pos_vel
-
-    W_p_Ci, _ = collidable_points_pos_vel(
-        model=model.physics_model,
-        q=data.state.physics_model.joint_positions,
-        qd=data.state.physics_model.joint_velocities,
-        xfb=data.state.physics_model.xfb(),
-    )
+    W_p_Ci = collidable_point_positions(model=model, data=data)
 
     terrain_height = jax.vmap(lambda x, y: model.terrain.height(x=x, y=y))(
-        W_p_Ci[0, :], W_p_Ci[1, :]
+        W_p_Ci[:, 0], W_p_Ci[:, 1]
     )
 
-    below_terrain = W_p_Ci[2, :] <= terrain_height
+    below_terrain = W_p_Ci[:, 2] <= terrain_height
 
     links_in_contact = jax.vmap(
         lambda link_index: jnp.where(
-            model.physics_model.gc.body == link_index,
+            jnp.array(model.kin_dyn_parameters.contact_parameters.body) == link_index,
             below_terrain,
             jnp.zeros_like(below_terrain, dtype=bool),
         ).any()
-    )(jnp.arange(model.number_of_links()))
+    )(js.link.names_to_idxs(link_names=link_names, model=model))
 
     return links_in_contact
 
 
 @jax.jit
 def estimate_good_soft_contacts_parameters(
-    model: Model.JaxSimModel,
+    model: js.model.JaxSimModel,
+    *,
+    standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
     static_friction_coefficient: jtp.FloatLike = 0.5,
     number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
     damping_ratio: jtp.FloatLike = 1.0,
     max_penetration: jtp.FloatLike | None = None,
-) -> soft_contacts.SoftContactsParams:
+) -> jaxsim.rbda.soft_contacts.SoftContactsParams:
     """
     Estimate good soft contacts parameters for the given model.
 
     Args:
         model: The model to consider.
+        standard_gravity: The standard gravity constant.
         static_friction_coefficient: The static friction coefficient.
         number_of_active_collidable_points_steady_state:
             The number of active collidable points in steady state supporting
@@ -160,16 +233,17 @@ def estimate_good_soft_contacts_parameters(
         specific application.
     """
 
-    def estimate_model_height(model: Model.JaxSimModel) -> jtp.Float:
+    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
         """"""
 
-        zero_data = Data.JaxSimModelData.build(
-            model=model, soft_contacts_params=soft_contacts.SoftContactsParams()
+        zero_data = js.data.JaxSimModelData.build(
+            model=model,
+            soft_contacts_params=jaxsim.rbda.soft_contacts.SoftContactsParams(),
         )
 
-        W_pz_CoM = Model.com_position(model=model, data=zero_data)[2]
+        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
 
-        if model.physics_model.is_floating_base:
+        if model.floating_base():
             W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
             return 2 * (W_pz_CoM - W_pz_C.min())
 
@@ -183,12 +257,15 @@ def estimate_good_soft_contacts_parameters(
 
     nc = number_of_active_collidable_points_steady_state
 
-    sc_parameters = soft_contacts.SoftContactsParams.build_default_from_physics_model(
-        physics_model=model.physics_model,
-        static_friction_coefficient=static_friction_coefficient,
-        max_penetration=max_δ,
-        number_of_active_collidable_points_steady_state=nc,
-        damping_ratio=damping_ratio,
+    sc_parameters = (
+        jaxsim.rbda.soft_contacts.SoftContactsParams.build_default_from_jaxsim_model(
+            model=model,
+            standard_gravity=standard_gravity,
+            static_friction_coefficient=static_friction_coefficient,
+            max_penetration=max_δ,
+            number_of_active_collidable_points_steady_state=nc,
+            damping_ratio=damping_ratio,
+        )
     )
 
     return sc_parameters