jaxsim 0.4.3.dev68__py3-none-any.whl → 0.4.3.dev77__py3-none-any.whl

jaxsim/terrain/terrain.py

@@ -46,66 +46,82 @@ class Terrain(abc.ABC):
 @jax_dataclasses.pytree_dataclass
 class FlatTerrain(Terrain):
 
-    z: float = dataclasses.field(default=0.0, kw_only=True)
+    _height: float = dataclasses.field(default=0.0, kw_only=True)
 
     @staticmethod
     def build(height: jtp.FloatLike) -> FlatTerrain:
 
-        return FlatTerrain(z=float(height))
+        return FlatTerrain(_height=float(height))
 
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
 
-        return jnp.array(self.z, dtype=float)
+        return jnp.array(self._height, dtype=float)
+
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+
+        return jnp.array([0.0, 0.0, 1.0], dtype=float)
 
     def __hash__(self) -> int:
 
-        return hash(self.z)
+        return hash(self._height)
 
     def __eq__(self, other: FlatTerrain) -> bool:
 
         if not isinstance(other, FlatTerrain):
             return False
 
-        return self.z == other.z
+        return self._height == other._height
 
 
 @jax_dataclasses.pytree_dataclass
 class PlaneTerrain(FlatTerrain):
 
-    plane_normal: tuple[float, float, float] = jax_dataclasses.field(
+    _normal: tuple[float, float, float] = jax_dataclasses.field(
         default=(0.0, 0.0, 1.0), kw_only=True
     )
 
     @staticmethod
-    def build(
-        plane_normal: jtp.VectorLike, plane_height_over_origin: jtp.FloatLike = 0.0
-    ) -> PlaneTerrain:
+    def build(height: jtp.FloatLike = 0.0, *, normal: jtp.VectorLike) -> PlaneTerrain:
         """
         Create a PlaneTerrain instance with a specified plane normal vector.
 
         Args:
-            plane_normal: The normal vector of the terrain plane.
-            plane_height_over_origin: The height of the plane over the origin.
+            normal: The normal vector of the terrain plane.
+            height: The height of the plane over the origin.
 
         Returns:
             PlaneTerrain: A PlaneTerrain instance.
         """
 
-        plane_normal = jnp.array(plane_normal, dtype=float)
-        plane_height_over_origin = jnp.array(plane_height_over_origin, dtype=float)
+        normal = jnp.array(normal, dtype=float)
+        height = jnp.array(height, dtype=float)
 
-        if plane_normal.shape != (3,):
+        if normal.shape != (3,):
             msg = "Expected a 3D vector for the plane normal, got '{}'."
-            raise ValueError(msg.format(plane_normal.shape))
+            raise ValueError(msg.format(normal.shape))
 
         # Make sure that the plane normal is a unit vector.
-        plane_normal = plane_normal / jnp.linalg.norm(plane_normal)
+        normal = normal / jnp.linalg.norm(normal)
 
         return PlaneTerrain(
-            z=float(plane_height_over_origin),
-            plane_normal=tuple(plane_normal.tolist()),
+            _height=height.item(),
+            _normal=tuple(normal.tolist()),
         )
 
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+        """
+        Compute the normal vector of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The normal vector of the terrain surface at the specified location.
+        """
+
+        return jnp.array(self._normal, dtype=float)
+
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
         """
         Compute the height of the terrain at a specific (x, y) location on a plane.
@@ -123,10 +139,10 @@ class PlaneTerrain(FlatTerrain):
         # The height over the origin: -D/C
 
         # Get the plane equation coefficients from the terrain normal.
-        A, B, C = self.plane_normal
+        A, B, C = self._normal
 
         # Compute the final coefficient D considering the terrain height.
-        D = -C * self.z
+        D = -C * self._height
 
         # Invert the plane equation to get the height at the given (x, y) coordinates.
         return jnp.array(-(A * x + B * y + D) / C).astype(float)
@@ -137,9 +153,9 @@ class PlaneTerrain(FlatTerrain):
 
         return hash(
             (
-                hash(self.z),
+                hash(self._height),
                 HashedNumpyArray.hash_of_array(
-                    array=jnp.array(self.plane_normal, dtype=float)
+                    array=jnp.array(self._normal, dtype=float)
                 ),
             )
         )
@@ -150,10 +166,10 @@ class PlaneTerrain(FlatTerrain):
             return False
 
         if not (
-            np.allclose(self.z, other.z)
+            np.allclose(self._height, other._height)
             and np.allclose(
-                np.array(self.plane_normal, dtype=float),
-                np.array(other.plane_normal, dtype=float),
+                np.array(self._normal, dtype=float),
+                np.array(other._normal, dtype=float),
             )
         ):
             return False

jaxsim/typing.py

@@ -16,7 +16,7 @@ Int = Scalar
 Bool = Scalar
 Float = Scalar
 
-PyTree = (
+PyTree: object = (
     dict[Hashable, TypeVar("PyTree")]
     | list[TypeVar("PyTree")]
     | tuple[TypeVar("PyTree")]

jaxsim/utils/jaxsim_dataclass.py

@@ -135,9 +135,10 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            leaf.shape if hasattr(leaf, "shape") else None
-            for leaf in jax.tree_util.tree_leaves(tree)
-            if hasattr(leaf, "shape")
+            map(
+                lambda leaf: getattr(leaf, "shape", None),
+                jax.tree_util.tree_leaves(tree),
+            )
         )
 
     @staticmethod
@@ -154,9 +155,10 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            leaf.dtype if hasattr(leaf, "dtype") else None
-            for leaf in jax.tree_util.tree_leaves(tree)
-            if hasattr(leaf, "dtype")
+            map(
+                lambda leaf: getattr(leaf, "dtype", None),
+                jax.tree_util.tree_leaves(tree),
+            )
         )
 
     @staticmethod
@@ -172,9 +174,10 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            leaf.weak_type if hasattr(leaf, "weak_type") else False
-            for leaf in jax.tree_util.tree_leaves(tree)
-            if hasattr(leaf, "weak_type")
+            map(
+                lambda leaf: getattr(leaf, "weak_type", None),
+                jax.tree_util.tree_leaves(tree),
+            )
         )
 
     @staticmethod

jaxsim/utils/wrappers.py

@@ -110,7 +110,7 @@ class HashedNumpyArray:
             return np.allclose(
                 self.array,
                 other.array,
-                **({dict(atol=self.precision)} if self.precision is not None else {}),
+                **(dict(atol=self.precision) if self.precision is not None else {}),
             )
 
         return hash(self) == hash(other)

{jaxsim-0.4.3.dev68.dist-info → jaxsim-0.4.3.dev77.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev68
+Version: 0.4.3.dev77
 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>
@@ -61,6 +61,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
 Requires-Dist: jax>=0.4.13
+Requires-Dist: jaxopt>=0.8.0
 Requires-Dist: jaxlib>=0.4.13
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax-dataclasses>=1.4.0

{jaxsim-0.4.3.dev68.dist-info → jaxsim-0.4.3.dev77.dist-info}/RECORD

@@ -1,29 +1,29 @@
-jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=XDf5LPSlhAhH48AO29kysLP_4FTR5VWOpS0LrK5RSfo,426
+jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
+jaxsim/_version.py,sha256=iooIC7hyVsbupWXsxWXV06oZTPyCALANUUOu7UZTm8Y,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
-jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
+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=HyEAjF7BySDDOlRahN0l7V15IPB0HPXuoM0twamuEW0,20913
-jaxsim/api/data.py,sha256=CUh9lvhVk3_clNQ26BUBGpjvFSsK_PrVWVMEWpMdHRM,27206
+jaxsim/api/contact.py,sha256=Ek1xSKB_VWjfqsqRYlK236ountKmGTl1M04cTYqHgsE,22142
+jaxsim/api/data.py,sha256=QldUHniJqKrdNtAcXuRaS9UyeslJ0Rjvb17UA0Ca5Tw,29008
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
-jaxsim/api/joint.py,sha256=L81bQe-noPT6_54KOSF7KBjRmEPAS433ULn2EcXI8vI,5115
+jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
 jaxsim/api/link.py,sha256=qPRtc8qqMRjZxUCZYXJMygbB6huDXBfIT1b1b8Durkw,18631
-jaxsim/api/model.py,sha256=HXoqCtQ3KStGoxhgvFm8P_Sc-lbEM4l5No2MoHzNlOk,65558
-jaxsim/api/ode.py,sha256=Vb2sN4zwpXnaJDD9-ziz2qvfmfa4jvIQ0fONbBIRGmU,13368
-jaxsim/api/ode_data.py,sha256=U7F6TL6bENAxpQQl4PupPoDG7d7VfTTFqDAs3xwu6Hs,20003
+jaxsim/api/model.py,sha256=TLjgacgTXm-2YRGDA0Id9pe9nxIem28KoAls6Tdk9WM,66241
+jaxsim/api/ode.py,sha256=ZshGdHptftku0yoUwBiBdd1iOqntH0vVEOjRHfL7Fao,13518
+jaxsim/api/ode_data.py,sha256=7RSoBhfCJdP6P9InQbDwdBVpClPMMuetewI-6AWm-_0,20276
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=ntjflaV3qWaFH_E65pAGZ6QipdnFsgQDasKtIKpxTe4,20432
+jaxsim/integrators/common.py,sha256=XIrJVJDO0ldaZ93WgoGNlFoRvazsRJTpO3DrK9kIXqM,20437
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
 jaxsim/integrators/variable_step.py,sha256=5StkFh9oQba34zlkIoXG2fUN78gbxkHePWbrpQ-QZOI,21274
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
-jaxsim/math/inertia.py,sha256=_hNpoeyEpAGr9ExDQJjckbjhk39luJFF-jv0SKqefnQ,1614
+jaxsim/math/inertia.py,sha256=01hz6wMFreN2jBA0rVoBS1YMVh77KvwuzXSOpI3pxNk,1614
 jaxsim/math/joint_model.py,sha256=EzAveaG5B6ZnCFNUzN30KEQUVesd83lfWXJarYR-kUw,9989
 jaxsim/math/quaternion.py,sha256=_WA7W3iv7px83sWO1V1n0-J78hqAlO4SL1-jofE-UZ4,4754
 jaxsim/math/rotation.py,sha256=k-nwT79zmWrys3NNAB-lGWxat7Kqm_6JnFRoimJ8rBg,2156
@@ -31,18 +31,18 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=KXzQgOnCfAtbXCwxhplpJ3F0JT3oEyeLVby1_uRAryQ,2892
 jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=XB-fgXuWMTFiaand5MZlLFQ5__Sh8MK5CJsxIU34MBk,25328
+jaxsim/mujoco/loaders.py,sha256=_8Af_5Yo0-lWHE-46BBMcrqSJnDNxr3peyc519DExtA,25322
 jaxsim/mujoco/model.py,sha256=AQksXemXWACJ3yvefV2G5HLwwBU9ISoJrOD1wlxdY5w,16386
 jaxsim/mujoco/visualizer.py,sha256=T1vU-w4NKSmgEkZ0FqVcGmIvYrYO0len2UBSsU4MOZ0,6978
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/parsers/kinematic_graph.py,sha256=KijMWKyhTLKSNUmOOk4sYQMgPh_OkA_brncL7gBRHaY,34757
+jaxsim/parsers/kinematic_graph.py,sha256=wT2bgaCS8VQJTHy2H9sENkVPDOiMkRikxEF1t_WaahQ,34748
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
 jaxsim/parsers/descriptions/collision.py,sha256=BQeIG-TKi4SVny23w6riDrQ5itC6VRwEMBX6HgAXHxA,3973
 jaxsim/parsers/descriptions/joint.py,sha256=VSb6C0FBBKMqwrHBKfc-Bbn4rl_J0RzUxMQlhIEvOPM,5185
 jaxsim/parsers/descriptions/link.py,sha256=Eh0W5qL7_Uw0GV-BkNKXhm9Q2dRTfIWCX5D-87zQkxA,3711
 jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwytayxds0,9833
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=HskeCqDsbtwH2BDk3vfxvx391wUTVGLaUXNvBrdNo-4,13486
+jaxsim/parsers/rod/parser.py,sha256=9EigYv2oGn4bfIY1q0Cd_55yVKfN2rXP_MuZSZqGxYM,13681
 jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5698
 jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
@@ -54,16 +54,17 @@ 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=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/rbda/contacts/common.py,sha256=VwAs742futAmLnDgbaOuLzNDBFiKDfYItdEZ4UcFgzE,2467
-jaxsim/rbda/contacts/rigid.py,sha256=8Vbnxng-ERZ5ka_eZGIBuhBDr2PNjc7m-Or255AfEw4,15862
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=0rrk73MKMVPkEXf1MTEmCnXaPWzzEw6U9RHZtLNIu_E,13646
+jaxsim/rbda/contacts/rigid.py,sha256=zbSM0miwpgC1rp1d0RoQ1q8pYiKdIkHV8iZimeEPC94,15153
 jaxsim/rbda/contacts/soft.py,sha256=_wvb5iZDjGcVg6rNQelN4LZN7qSC2NIp0HdKvZmlGfk,15647
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
-jaxsim/terrain/terrain.py,sha256=ctyNANIFSM3tZmamprjaEDcWgUSP0oNJbmT1zw9RjPs,4565
+jaxsim/terrain/terrain.py,sha256=xUQg47yGxIOcTkLPbnO3sruEGBhoCd16j1evTGlmNjI,5010
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
-jaxsim/utils/jaxsim_dataclass.py,sha256=5xJbY0G8d7C0OTNIW9T4vQxiDak6TGZT9gpNOvRykFI,11373
+jaxsim/utils/jaxsim_dataclass.py,sha256=FSiUvdnq4Y1T9Jaa_mw4ZBQJe8H7deLr3Kupxtlh4iI,11322
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim/utils/wrappers.py,sha256=JhLUh1g8iU-lhjbuZRfkscPZhYlLCOorVM2Xl3ulRBI,4054
-jaxsim-0.4.3.dev68.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev68.dist-info/METADATA,sha256=IrZMXHUptvvLA5YgloveNIge4OdEBjT-DxhdHBrn_WM,17247
-jaxsim-0.4.3.dev68.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-jaxsim-0.4.3.dev68.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev68.dist-info/RECORD,,
+jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
+jaxsim-0.4.3.dev77.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev77.dist-info/METADATA,sha256=SfRw58102asAL9OAPXqsvh9nr8FVvLamf9dCNL9XXPo,17276
+jaxsim-0.4.3.dev77.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev77.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev77.dist-info/RECORD,,

{jaxsim-0.4.3.dev68.dist-info → jaxsim-0.4.3.dev77.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any