jaxsim 0.4.3.dev143__py3-none-any.whl → 0.4.3.dev159__py3-none-any.whl

jaxsim/rbda/contacts/soft.py

@@ -10,10 +10,11 @@ import jax_dataclasses
 import jaxsim.api as js
 import jaxsim.math
 import jaxsim.typing as jtp
+from jaxsim import logging
 from jaxsim.math import StandardGravity
 from jaxsim.terrain import FlatTerrain, Terrain
 
-from .common import ContactModel, ContactsParams, ContactsState
+from .common import ContactModel, ContactsParams
 
 try:
     from typing import Self
@@ -192,13 +193,67 @@ class SoftContacts(ContactModel):
     """Soft contacts model."""
 
     parameters: SoftContactsParams = dataclasses.field(
-        default_factory=SoftContactsParams
+        default_factory=SoftContactsParams.build
     )
 
     terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
+        default_factory=FlatTerrain.build
     )
 
+    @classmethod
+    def build(
+        cls: type[Self],
+        parameters: SoftContactsParams | None = None,
+        terrain: Terrain | None = None,
+        model: js.model.JaxSimModel | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `SoftContacts` instance with specified parameters.
+
+        Args:
+            parameters: The parameters of the soft contacts model.
+            terrain: The considered terrain.
+            model:
+                The robot model considered by the contact model.
+                If passed, it is used to estimate good default parameters.
+
+        Returns:
+            The `SoftContacts` instance.
+        """
+
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
+
+        # Build the contact parameters if not provided. Use the model to estimate
+        # good default parameters, if passed. Users can later override these default
+        # parameters with their own values -- possibly tuned better.
+        if parameters is None:
+            parameters = (
+                SoftContactsParams.build_default_from_jaxsim_model(model=model)
+                if model is not None
+                else cls.__dataclass_fields__["parameters"].default_factory()
+            )
+
+        return SoftContacts(
+            parameters=parameters,
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+        )
+
+    @classmethod
+    def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
+        """
+        Build zero state variables of the contact model.
+        """
+
+        # Initialize the material deformation to zero.
+        tangential_deformation = jnp.zeros(
+            shape=(len(model.kin_dyn_parameters.contact_parameters.body), 3),
+            dtype=float,
+        )
+
+        return {"tangential_deformation": tangential_deformation}
+
     @staticmethod
     @functools.partial(jax.jit, static_argnames=("terrain",))
     def hunt_crossley_contact_model(
@@ -380,8 +435,7 @@ class SoftContacts(ContactModel):
         W_p_C, W_ṗ_C = js.contact.collidable_point_kinematics(model=model, data=data)
 
         # Extract the material deformation corresponding to the collidable points.
-        assert isinstance(data.state.contact, SoftContactsState)
-        m = data.state.contact.tangential_deformation
+        m = data.state.extended["tangential_deformation"]
 
         # Compute the contact forces for all collidable points.
         # Since we treat them as independent, we can vmap the computation.
@@ -423,131 +477,3 @@ class SoftContacts(ContactModel):
         δ̇ = jnp.where(δ > 0, δ̇, 0.0)
 
         return δ, δ̇, n̂
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContactsState(ContactsState):
-    """
-    Class storing the state of the soft contacts model.
-
-    Attributes:
-        tangential_deformation:
-            The matrix of 3D tangential material deformations corresponding to
-            each collidable point.
-    """
-
-    tangential_deformation: jtp.Matrix
-
-    def __hash__(self) -> int:
-
-        return hash(
-            tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
-        )
-
-    def __eq__(self: Self, other: Self) -> bool:
-
-        if not isinstance(other, type(self)):
-            return False
-
-        return hash(self) == hash(other)
-
-    @classmethod
-    def build_from_jaxsim_model(
-        cls: type[Self],
-        model: js.model.JaxSimModel | None = None,
-        tangential_deformation: jtp.MatrixLike | None = None,
-    ) -> Self:
-        """
-        Build a `SoftContactsState` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the soft contacts state.
-            tangential_deformation: The matrix of 3D tangential material deformations.
-
-        Returns:
-            The `SoftContactsState` built from the `JaxSimModel`.
-
-        Note:
-            If any of the state components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return cls.build(
-            tangential_deformation=tangential_deformation,
-            number_of_collidable_points=len(
-                model.kin_dyn_parameters.contact_parameters.body
-            ),
-        )
-
-    @classmethod
-    def build(
-        cls: type[Self],
-        *,
-        tangential_deformation: jtp.MatrixLike | None = None,
-        number_of_collidable_points: int | None = None,
-    ) -> Self:
-        """
-        Create a `SoftContactsState`.
-
-        Args:
-            tangential_deformation:
-                The matrix of 3D tangential material deformations corresponding to
-                each collidable point.
-            number_of_collidable_points: The number of collidable points.
-
-        Returns:
-            A `SoftContactsState` instance.
-        """
-
-        tangential_deformation = (
-            jnp.atleast_2d(tangential_deformation)
-            if tangential_deformation is not None
-            else jnp.zeros(shape=(number_of_collidable_points, 3))
-        ).astype(float)
-
-        if tangential_deformation.shape[1] != 3:
-            raise RuntimeError("The tangential deformation matrix must have 3 columns.")
-
-        if (
-            number_of_collidable_points is not None
-            and tangential_deformation.shape[0] != number_of_collidable_points
-        ):
-            msg = "The number of collidable points must match the number of rows "
-            msg += "in the tangential deformation matrix."
-            raise RuntimeError(msg)
-
-        return cls(tangential_deformation=tangential_deformation)
-
-    @classmethod
-    def zero(cls: type[Self], *, model: js.model.JaxSimModel) -> Self:
-        """
-        Build a zero `SoftContactsState` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the soft contacts state.
-
-        Returns:
-            A zero `SoftContactsState` instance.
-        """
-
-        return cls.build_from_jaxsim_model(model=model)
-
-    def valid(self, *, model: js.model.JaxSimModel) -> jtp.BoolLike:
-        """
-        Check if the `SoftContactsState` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `SoftContactsState` against.
-
-        Returns:
-            `True` if the soft contacts state is valid for the given `JaxSimModel`,
-            `False` otherwise.
-        """
-
-        shape = self.tangential_deformation.shape
-        expected = (len(model.kin_dyn_parameters.contact_parameters.body), 3)
-
-        if shape != expected:
-            return False
-
-        return True

jaxsim/terrain/terrain.py

@@ -49,7 +49,7 @@ class FlatTerrain(Terrain):
     _height: float = dataclasses.field(default=0.0, kw_only=True)
 
     @staticmethod
-    def build(height: jtp.FloatLike) -> FlatTerrain:
+    def build(height: jtp.FloatLike = 0.0) -> FlatTerrain:
 
         return FlatTerrain(_height=float(height))
 

{jaxsim-0.4.3.dev143.dist-info → jaxsim-0.4.3.dev159.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev143
+Version: 0.4.3.dev159
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev143.dist-info → jaxsim-0.4.3.dev159.dist-info}/RECORD

@@ -1,20 +1,20 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=AUMtPwp-9k4PBssfGUg1S2TuNzVQkbYfnNIjR31wCZY,428
+jaxsim/_version.py,sha256=p1_cPvBnSIuAbg5RWked461etxKJN-ANdqjllGN6J4s,428
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
-jaxsim/api/contact.py,sha256=ktol-c2UHDqt3Hd6NpMhK86NlyadIyB5eY7chA5GhkY,22568
-jaxsim/api/data.py,sha256=bSuBVKKksYtANLnxCkMQ3u2JGbn5mud7g8G5aLZawls,28988
+jaxsim/api/contact.py,sha256=RYIaVqDMizxfretC72Ye5vfr8208UsFyzqGE-C1oxxA,22204
+jaxsim/api/data.py,sha256=kDP1s5mm8F-LpIqKrdaCmHhJEMyigFDyuaAXdMukhWY,29588
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=ElahFk_RCcLvjTidH2qDOsY-m1gN1hXitCv4SvfgGYY,29260
 jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=Y0A9p7uKpfBdZLC9o_pTw7WQrmAaNJJBK1G97YkDRjQ,68260
-jaxsim/api/ode.py,sha256=B9CANQGLbrKAkmTxvWvkb0UGKC__r1nlHf-bGw963q0,13999
-jaxsim/api/ode_data.py,sha256=6q3IECaA351hLSGRovrOSKLo_P8IhJrW-p2-ZDlzGG8,20436
+jaxsim/api/model.py,sha256=IcJhFdVjswSncurVpQXoa7uXVqSsvAuZVs-DXDsXi08,69021
+jaxsim/api/ode.py,sha256=_rTgDVNLQWbZ-E8ApKbq4JQ6ctCcQeDvCzAf6d35hY8,14061
+jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
 jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=78MBs89GxsL0wU2yAexjvBZt3HEtfZoGVIN9f0a8yTc,20305
@@ -52,19 +52,19 @@ jaxsim/rbda/forward_kinematics.py,sha256=2GmEoWsrioVl_SAbKRKfhOLz57pY4aR81PKRdul
 jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
 jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
-jaxsim/rbda/contacts/__init__.py,sha256=0UnO9ZR3BwdjQa276jOFbPi90pporr32LSc0qa9UUm4,369
-jaxsim/rbda/contacts/common.py,sha256=-eM8d1kvJ2E_2_kAgZJk4s3x8vDZHNSyOAinwPmRmEk,3469
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=G3NoFrIGqzr_hrScLtsbmliTSqcVxDxheA1w2HRcU7A,15794
-jaxsim/rbda/contacts/rigid.py,sha256=WzPPo_UACOo_7olv_R1PelRjdjGqdWHkQ6hkSSfdBZk,15689
-jaxsim/rbda/contacts/soft.py,sha256=NzzCYw5rvK8Fx_qH3fiMzPgey-KoxmRe9xkF3fluidE,18866
+jaxsim/rbda/contacts/__init__.py,sha256=Vqtk2GeSAXh6YxiqadY-fe4PY__IUYiOwnLgSdpLqW0,271
+jaxsim/rbda/contacts/common.py,sha256=_yrxTM16Je9ck5aM95ndk8Kwu_oijxG9Jaf1jEjHEYw,4332
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=Ob5LdKe3D7tGlIdT4LamJ6_F0j5pzUmWNYoWqy8Di98,17169
+jaxsim/rbda/contacts/rigid.py,sha256=1TTiGXSOipO8l5FDTtxqRNo1ArCNtDg-Yr3olPgBLGs,17588
+jaxsim/rbda/contacts/soft.py,sha256=TMCUDtFmNIae04LCla57iXMjdt9F5qTFjYEnP5NdLFg,16809
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
-jaxsim/terrain/terrain.py,sha256=Y0TGnUAGPuaeeSN8vbaSFjMXP5GYy3nxMCETjpUIMSA,5009
+jaxsim/terrain/terrain.py,sha256=K91HEzPqTSyNrc_j1KfAAEF_5oDeuk_-jnnZGrcMEcY,5015
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev143.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev143.dist-info/METADATA,sha256=jZMIz7kCovjDtiHxhrp18oWpaYUhdX8t056wtbTv1us,17276
-jaxsim-0.4.3.dev143.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev143.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev143.dist-info/RECORD,,
+jaxsim-0.4.3.dev159.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev159.dist-info/METADATA,sha256=3AP-0_JvmN0Co6aNIBawlZVlTfi1nMi4V-XSchb91lM,17276
+jaxsim-0.4.3.dev159.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev159.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev159.dist-info/RECORD,,