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

jaxsim/utils/jaxsim_dataclass.py

@@ -2,7 +2,9 @@ import abc
 import contextlib
 import copy
 import dataclasses
-from typing import ClassVar, Generator
+import functools
+from collections.abc import Iterator
+from typing import Any, Callable, ClassVar, Sequence, Type
 
 import jax.flatten_util
 import jax_dataclasses
@@ -19,51 +21,219 @@ except ImportError:
 
 @jax_dataclasses.pytree_dataclass
 class JaxsimDataclass(abc.ABC):
-    """"""
+    """Class extending `jax_dataclasses.pytree_dataclass` instances with utilities."""
 
     # This attribute is set by jax_dataclasses
     __mutability__: ClassVar[Mutability] = Mutability.FROZEN
 
     @contextlib.contextmanager
-    def editable(self: Self, validate: bool = True) -> Generator[Self, None, None]:
-        """"""
+    def editable(self: Self, validate: bool = True) -> Iterator[Self]:
+        """
+        Context manager to operate on a mutable copy of the object.
+
+        Args:
+            validate: Whether to validate the output PyTree upon exiting the context.
+
+        Yields:
+            A mutable copy of the object.
+
+        Note:
+            This context manager is useful to operate on an r/w copy of a PyTree making
+            sure that the output object does not trigger JIT recompilations.
+        """
 
         mutability = (
             Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
         )
 
-        with JaxsimDataclass.mutable_context(self.copy(), mutability=mutability) as obj:
+        with self.copy().mutable_context(mutability=mutability) as obj:
             yield obj
 
     @contextlib.contextmanager
     def mutable_context(
         self: Self, mutability: Mutability, restore_after_exception: bool = True
-    ) -> Generator[Self, None, None]:
-        """"""
+    ) -> Iterator[Self]:
+        """
+        Context manager to temporarily change the mutability of the object.
+
+        Args:
+            mutability: The mutability to set.
+            restore_after_exception:
+                Whether to restore the original object in case of an exception
+                occurring within the context.
+
+        Yields:
+            The object with the new mutability.
+
+        Note:
+            This context manager is useful to operate in place on a PyTree without
+            the need to make a copy while optionally keeping active the checks on
+            the PyTree structure, shapes, and dtypes.
+        """
 
         if restore_after_exception:
             self_copy = self.copy()
 
-        original_mutability = self._mutability()
+        original_mutability = self.mutability()
+
+        original_dtypes = JaxsimDataclass.get_leaf_dtypes(tree=self)
+        original_shapes = JaxsimDataclass.get_leaf_shapes(tree=self)
+        original_weak_types = JaxsimDataclass.get_leaf_weak_types(tree=self)
+        original_structure = jax.tree_util.tree_structure(tree=self)
 
-        def restore_self():
-            self._set_mutability(mutability=Mutability.MUTABLE)
+        def restore_self() -> None:
+            self.set_mutability(mutability=Mutability.MUTABLE_NO_VALIDATION)
             for f in dataclasses.fields(self_copy):
                 setattr(self, f.name, getattr(self_copy, f.name))
 
         try:
-            self._set_mutability(mutability)
+            self.set_mutability(mutability)
             yield self
+
+            if mutability is not Mutability.MUTABLE_NO_VALIDATION:
+                new_structure = jax.tree_util.tree_structure(tree=self)
+                if original_structure != new_structure:
+                    msg = "Pytree structure has changed from {} to {}"
+                    raise ValueError(msg.format(original_structure, new_structure))
+
+                new_shapes = JaxsimDataclass.get_leaf_shapes(tree=self)
+                if original_shapes != new_shapes:
+                    msg = "Leaves shapes have changed from {} to {}"
+                    raise ValueError(msg.format(original_shapes, new_shapes))
+
+                new_dtypes = JaxsimDataclass.get_leaf_dtypes(tree=self)
+                if original_dtypes != new_dtypes:
+                    msg = "Leaves dtypes have changed from {} to {}"
+                    raise ValueError(msg.format(original_dtypes, new_dtypes))
+
+                new_weak_types = JaxsimDataclass.get_leaf_weak_types(tree=self)
+                if original_weak_types != new_weak_types:
+                    msg = "Leaves weak types have changed from {} to {}"
+                    raise ValueError(msg.format(original_weak_types, new_weak_types))
+
         except Exception as e:
             if restore_after_exception:
                 restore_self()
-            self._set_mutability(original_mutability)
+            self.set_mutability(original_mutability)
             raise e
+
         finally:
-            self._set_mutability(original_mutability)
+            self.set_mutability(original_mutability)
+
+    @staticmethod
+    def get_leaf_shapes(tree: jtp.PyTree) -> tuple[tuple[int, ...] | None]:
+        """
+        Helper method to get the leaf shapes of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple containing the leaf shapes of the PyTree or `None` is the leaf is
+            not a numpy-like array.
+        """
+
+        return tuple(  # noqa
+            leaf.shape if hasattr(leaf, "shape") else None
+            for leaf in jax.tree_util.tree_leaves(tree)
+            if hasattr(leaf, "shape")
+        )
+
+    @staticmethod
+    def get_leaf_dtypes(tree: jtp.PyTree) -> tuple:
+        """
+        Helper method to get the leaf dtypes of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple containing the leaf dtypes of the PyTree or `None` is the leaf is
+            not a numpy-like array.
+        """
+
+        return tuple(
+            leaf.dtype if hasattr(leaf, "dtype") else None
+            for leaf in jax.tree_util.tree_leaves(tree)
+            if hasattr(leaf, "dtype")
+        )
+
+    @staticmethod
+    def get_leaf_weak_types(tree: jtp.PyTree) -> tuple[bool, ...]:
+        """
+        Helper method to get the leaf weak types of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple marking whether the leaf contains a JAX array with weak type.
+        """
+
+        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")
+        )
+
+    @staticmethod
+    def check_compatibility(*trees: Sequence[Any]) -> None:
+        """
+        Check whether the PyTrees are compatible in structure, shape, and dtype.
+
+        Args:
+            *trees: The PyTrees to compare.
+
+        Raises:
+            ValueError: If the PyTrees have incompatible structures, shapes, or dtypes.
+        """
+
+        target_structure = jax.tree_util.tree_structure(trees[0])
+
+        compatible_structure = functools.reduce(
+            lambda compatible, tree: compatible
+            and jax.tree_util.tree_structure(tree) == target_structure,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_structure:
+            raise ValueError("Pytrees have incompatible structures.")
+
+        target_shapes = JaxsimDataclass.get_leaf_shapes(trees[0])
+
+        compatible_shapes = functools.reduce(
+            lambda compatible, tree: compatible
+            and JaxsimDataclass.get_leaf_shapes(tree) == target_shapes,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_shapes:
+            raise ValueError("Pytrees have incompatible shapes.")
+
+        target_dtypes = JaxsimDataclass.get_leaf_dtypes(trees[0])
+
+        compatible_dtypes = functools.reduce(
+            lambda compatible, tree: compatible
+            and JaxsimDataclass.get_leaf_dtypes(tree) == target_dtypes,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_dtypes:
+            raise ValueError("Pytrees have incompatible dtypes.")
 
     def is_mutable(self, validate: bool = False) -> bool:
-        """"""
+        """
+        Check whether the object is mutable.
+
+        Args:
+            validate: Additionally checks if the object also has validation enabled.
+
+        Returns:
+            True if the object is mutable, False otherwise.
+        """
 
         return (
             self.__mutability__ is Mutability.MUTABLE
@@ -71,39 +241,120 @@ class JaxsimDataclass(abc.ABC):
             else self.__mutability__ is Mutability.MUTABLE_NO_VALIDATION
         )
 
-    def set_mutability(self, mutable: bool = True, validate: bool = False) -> None:
-        if not mutable:
-            mutability = Mutability.FROZEN
-        else:
-            mutability = (
-                Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
-            )
+    def mutability(self) -> Mutability:
+        """
+        Get the mutability type of the object.
 
-        self._set_mutability(mutability=mutability)
+        Returns:
+            The mutability type of the object.
+        """
 
-    def _mutability(self) -> Mutability:
         return self.__mutability__
 
-    def _set_mutability(self, mutability: Mutability) -> None:
+    def set_mutability(self, mutability: Mutability) -> None:
+        """
+        Set the mutability of the object in-place.
+
+        Args:
+            mutability: The desired mutability type.
+        """
+
         jax_dataclasses._copy_and_mutate._mark_mutable(
             self, mutable=mutability, visited=set()
         )
 
     def mutable(self: Self, mutable: bool = True, validate: bool = False) -> Self:
-        self.set_mutability(mutable=mutable, validate=validate)
+        """
+        Return a mutable reference of the object.
+
+        Args:
+            mutable: Whether to make the object mutable.
+            validate: Whether to enable validation on the object.
+
+        Returns:
+            A mutable reference of the object.
+        """
+
+        if mutable:
+            mutability = (
+                Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
+            )
+        else:
+            mutability = Mutability.FROZEN
+
+        self.set_mutability(mutability=mutability)
         return self
 
     def copy(self: Self) -> Self:
+        """
+        Return a copy of the object.
+
+        Returns:
+            A copy of the object.
+        """
+
+        # Make a copy calling tree_map.
         obj = jax.tree_util.tree_map(lambda leaf: leaf, self)
-        obj._set_mutability(mutability=self._mutability())
+
+        # Make sure that the copied object and all the copied leaves have the same
+        # mutability of the original object.
+        obj.set_mutability(mutability=self.mutability())
+
         return obj
 
     def replace(self: Self, validate: bool = True, **kwargs) -> Self:
-        with self.editable(validate=validate) as obj:
-            _ = [obj.__setattr__(k, copy.copy(v)) for k, v in kwargs.items()]
+        """
+        Return a new object replacing in-place the specified fields with new values.
+
+        Args:
+            validate: Whether to validate that the new fields do not alter the PyTree.
+            **kwargs: The fields to replace.
+
+        Returns:
+            A reference of the object with the specified fields replaced.
+        """
+
+        # Use the dataclasses replace method.
+        obj = dataclasses.replace(self, **kwargs)
+
+        if validate:
+            JaxsimDataclass.check_compatibility(self, obj)
+
+        # Make sure that all the new leaves have the same mutability of the object.
+        obj.set_mutability(mutability=self.mutability())
 
-        obj._set_mutability(mutability=self._mutability())
         return obj
 
     def flatten(self) -> jtp.VectorJax:
-        return jax.flatten_util.ravel_pytree(self)[0]
+        """
+        Flatten the object into a 1D vector.
+
+        Returns:
+            A 1D vector containing the flattened object.
+        """
+
+        return self.flatten_fn()(self)
+
+    @classmethod
+    def flatten_fn(cls: Type[Self]) -> Callable[[Self], jtp.VectorJax]:
+        """
+        Return a function to flatten the object into a 1D vector.
+
+        Returns:
+            A function to flatten the object into a 1D vector.
+        """
+
+        return lambda pytree: jax.flatten_util.ravel_pytree(pytree)[0]
+
+    def unflatten_fn(self: Self) -> Callable[[jtp.VectorJax], Self]:
+        """
+        Return a function to unflatten a 1D vector into the object.
+
+        Returns:
+            A function to unflatten a 1D vector into the object.
+
+        Notes:
+            Due to JAX internals, the function to unflatten a PyTree needs to be
+            created from an existing instance of the PyTree.
+        """
+        return jax.flatten_util.ravel_pytree(self)[1]

{jaxsim-0.2.dev191.dist-info → jaxsim-0.2.dev364.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev191
+Version: 0.2.dev364
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
@@ -31,12 +31,12 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
-Requires-Dist: jax <0.4.25,>=0.4.13
-Requires-Dist: jaxlib <0.4.25,>=0.4.13
+Requires-Dist: jax >=0.4.13
+Requires-Dist: jaxlib >=0.4.13
 Requires-Dist: jaxlie >=1.3.0
 Requires-Dist: jax-dataclasses >=1.4.0
 Requires-Dist: pptree
-Requires-Dist: rod
+Requires-Dist: rod >=0.2.0
 Requires-Dist: typing-extensions ; python_version < "3.12"
 Provides-Extra: all
 Requires-Dist: black[jupyter] ; extra == 'all'
@@ -44,7 +44,6 @@ Requires-Dist: isort ; extra == 'all'
 Requires-Dist: pre-commit ; extra == 'all'
 Requires-Dist: idyntree ; extra == 'all'
 Requires-Dist: pytest >=6.0 ; extra == 'all'
-Requires-Dist: pytest-forked ; extra == 'all'
 Requires-Dist: pytest-icdiff ; extra == 'all'
 Requires-Dist: robot-descriptions ; extra == 'all'
 Requires-Dist: lxml ; extra == 'all'
@@ -57,7 +56,6 @@ Requires-Dist: pre-commit ; extra == 'style'
 Provides-Extra: testing
 Requires-Dist: idyntree ; extra == 'testing'
 Requires-Dist: pytest >=6.0 ; extra == 'testing'
-Requires-Dist: pytest-forked ; extra == 'testing'
 Requires-Dist: pytest-icdiff ; extra == 'testing'
 Requires-Dist: robot-descriptions ; extra == 'testing'
 Provides-Extra: viz

jaxsim-0.2.dev364.dist-info/RECORD

@@ -0,0 +1,64 @@
+jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
+jaxsim/_version.py,sha256=yZse2Rzb5HFwkUDJRRD3kYlkmiLmp2arTSwcTm3cl1E,423
+jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
+jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
+jaxsim/api/__init__.py,sha256=fNTCPUeDfOAizRd4RsW3Epv0sLTu0KJGoFRSEsi75VM,162
+jaxsim/api/com.py,sha256=Qtm_6qpiK4WtDVn6JMUHa8DySgBl9CjgKCybJqZ58Lc,7379
+jaxsim/api/common.py,sha256=6oqZO-QTYr2mpMx5qRrAnCIjQpjIJVe7MlavNFIKbNA,6638
+jaxsim/api/contact.py,sha256=Ve4ZOWkLEBRgK3KhtICxKY7YzsxYvc3lO-pPRBjqSnY,8659
+jaxsim/api/data.py,sha256=1AJyKjmXdsWEf_CkvOXRmvRsDZamG706mAAg1Gttll0,26773
+jaxsim/api/joint.py,sha256=q31Kp3Cqv-yTcxijjzbj_QADFnGQyjb2al9fYZtzedo,4763
+jaxsim/api/kin_dyn_parameters.py,sha256=G4mtSi8fElYe0yttLgsxSOPf7vcK-yqTu06Aa5SSrYg,26012
+jaxsim/api/link.py,sha256=LZVcQhQsTKsfR13KewFtEMYu4siVJl7mqoDwYsoFFes,9240
+jaxsim/api/model.py,sha256=mFdEwVuzIR7Lvj4oiIsA1n1oxpRZXWyB0IDzIxcG33Q,43689
+jaxsim/api/ode.py,sha256=rbSruK0Dkp09oBgHbB_-NrZS4o2tY9geK0yLLJfXzpM,9821
+jaxsim/api/ode_data.py,sha256=dwRFVPJ30XMmdUbPXEu7YxsQ97jZP4L4fd5ZzhrO5Ys,22184
+jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
+jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
+jaxsim/integrators/common.py,sha256=APmQVXKdN9JMIq7wrVKq2HlI1UKpqvLqaTqAq0VfJJ0,20700
+jaxsim/integrators/fixed_step.py,sha256=JXaEyEzfSiYea0GnPA7l27J3X0YPB0e25D4qfrxAvzQ,2766
+jaxsim/integrators/variable_step.py,sha256=jq3PStzFiMciU7lux6CTj4B3gVOfSpYgK2oz2yzIbdo,21380
+jaxsim/math/__init__.py,sha256=inJ9nRFkqstuGa8OyFkfWVudo5U9Ug4WgDBuKva8AIA,337
+jaxsim/math/adjoint.py,sha256=DT21izjVW497GRrgNfx8tv0ZeWW5QncWMGMhI0acUNw,4425
+jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
+jaxsim/math/inertia.py,sha256=UAB7ym4gXFanejcs_ovZMpteHCc6poWYmt-mLmd5hhk,1640
+jaxsim/math/joint_model.py,sha256=LKLB26VMz6vx9JLdFUWhGyrElYFEQV-bJiQO5kaZUGY,10896
+jaxsim/math/quaternion.py,sha256=X9b8jHf0QemKUjIZSnXRJc3DdMr42CBhBy_mi9_X_AM,5068
+jaxsim/math/rotation.py,sha256=Z90daUjGpuNEVLfWB3SVtM9EtwAIaneVj9A9UpWXqhA,2182
+jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
+jaxsim/math/transform.py,sha256=nqH6ofde6VWjfHihmYdXvDxKFChpOPH6AsoqkUI1Og0,2928
+jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
+jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
+jaxsim/mujoco/loaders.py,sha256=8sXc_tsDFWBYl8nesgFarYd3hA-PESLMrXsnR3Siz1Y,16400
+jaxsim/mujoco/model.py,sha256=0kG2GERxjVFqWZ1K3352rgUNfchB4kRtIrsvv4pS4oc,10766
+jaxsim/mujoco/visualizer.py,sha256=-qg26t5tleTva6zzQmc5SdnlC8XZ1ZAwZ_lDjdwHJ0A,4400
+jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
+jaxsim/parsers/kinematic_graph.py,sha256=2B5gtUboiSVJIm6PegbwHb5g_iXltG0R9_7h8RJ-92M,23785
+jaxsim/parsers/descriptions/__init__.py,sha256=EbTfnrK3oCxA3pNv--YUwllJ6uICENvFgAdRbYtS9ts,238
+jaxsim/parsers/descriptions/collision.py,sha256=HUWwuRgI9KznY29FFw1_zU3bGigDEezrcPOJSxSJGNU,3382
+jaxsim/parsers/descriptions/joint.py,sha256=hpH0ANvIhbEQk-NGRmWIvPv3lXW385TBIMWNgz5rzM4,4106
+jaxsim/parsers/descriptions/link.py,sha256=hqLLitrAXnouia6ULk1BPOIEfRxrXwHmoPsi306IZW8,2859
+jaxsim/parsers/descriptions/model.py,sha256=wenuDrjoBf6prkzm9WyYT0nFWc0l6WBpKNjLoRUDPxo,8937
+jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
+jaxsim/parsers/rod/parser.py,sha256=mFi1baSJte6EMmWLpVjVuCpicfAAF48aFUzoKzYzPpo,12555
+jaxsim/parsers/rod/utils.py,sha256=xoxNdb4IgOePFFnblNo6UrgYCRm53mB3No3yXM6qlpw,6471
+jaxsim/rbda/__init__.py,sha256=HLwxeU-IxaRpFGUCSQv-LDv20JHTt3Xj7ELiRbRieS8,319
+jaxsim/rbda/aba.py,sha256=0OoCzHhf1v-qqr1y5PIrD7_mPwAlid0fjXxUrIa5E_s,9118
+jaxsim/rbda/collidable_points.py,sha256=4ZNJbEj2nEi15jBLR-GNbdaqKgkN58FBgqd_TXupEgg,4948
+jaxsim/rbda/crba.py,sha256=GodskOZjtrSlbQAqxRv1un_706O7BaJK-U2qa18vJk8,4741
+jaxsim/rbda/forward_kinematics.py,sha256=OHugNU7C0UxYAW0o1rqH1ZgniSwurz6L1T1MJxfxq08,3418
+jaxsim/rbda/jacobian.py,sha256=9LGGy9ya5m5U0mBmV1NFH5XYZpEMYbx74qnYBvZs7Ok,6360
+jaxsim/rbda/rnea.py,sha256=DjwkvXQVUSUclM3Uy3UPZ2tao91R5dGd4o7TsS2qObI,7650
+jaxsim/rbda/soft_contacts.py,sha256=2EZ9Lw4nFWqXTMEeYsirl17H61s82SmTZllKVsP1Yek,10759
+jaxsim/rbda/utils.py,sha256=zpbFM2Iq8cntku0BFVu9nfEqZhInCWi9D2INT6MFEI8,5003
+jaxsim/terrain/__init__.py,sha256=dzekq9yyj3DKTsCARteqc81lAw3OSnl6EhXn8_Q_ozI,64
+jaxsim/terrain/terrain.py,sha256=q0xkWqEShVq-p1j2abTLZq8sEhjyJwquxQKm80PaHhM,2161
+jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
+jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
+jaxsim/utils/jaxsim_dataclass.py,sha256=9FtzgXQPSa6AN26FleDgXGsJMse4dtVACDESrn0g3bw,11359
+jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
+jaxsim-0.2.dev364.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
+jaxsim-0.2.dev364.dist-info/METADATA,sha256=8p533sQ7hCyjA28HWGZ3NvohEIc56BSi06HlEVs6tJA,7630
+jaxsim-0.2.dev364.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.dev364.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.dev364.dist-info/RECORD,,

{jaxsim-0.2.dev191.dist-info → jaxsim-0.2.dev364.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

jaxsim/high_level/__init__.py

@@ -1,2 +0,0 @@
-from . import common, joint, link, model
-from .common import VelRepr

jaxsim/high_level/common.py

@@ -1,11 +0,0 @@
-import enum
-
-
-class VelRepr(enum.IntEnum):
-    """
-    Enumeration of all supported 6D velocity representations.
-    """
-
-    Body = enum.auto()
-    Mixed = enum.auto()
-    Inertial = enum.auto()

jaxsim/high_level/joint.py

@@ -1,148 +0,0 @@
-import dataclasses
-import functools
-from typing import Any
-
-import jax.numpy as jnp
-import jax_dataclasses
-from jax_dataclasses import Static
-
-import jaxsim.parsers
-import jaxsim.typing as jtp
-from jaxsim.utils import Vmappable, not_tracing, oop
-
-
-@jax_dataclasses.pytree_dataclass
-class Joint(Vmappable):
-    """
-    High-level class to operate in r/o on a single joint of a simulated model.
-    """
-
-    joint_description: Static[jaxsim.parsers.descriptions.JointDescription]
-
-    _parent_model: Any = dataclasses.field(
-        default=None, repr=False, compare=False, hash=False
-    )
-
-    @property
-    def parent_model(self) -> "jaxsim.high_level.model.Model":
-        """"""
-
-        return self._parent_model
-
-    @functools.partial(oop.jax_tf.method_ro, jit=False)
-    def valid(self) -> jtp.Bool:
-        """"""
-
-        return jnp.array(self.parent_model is not None, dtype=bool)
-
-    @functools.partial(oop.jax_tf.method_ro, jit=False)
-    def index(self) -> jtp.Int:
-        """"""
-
-        return jnp.array(self.joint_description.index, dtype=int)
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def dofs(self) -> jtp.Int:
-        """"""
-
-        return jnp.array(1, dtype=int)
-
-    @functools.partial(oop.jax_tf.method_ro, jit=False, vmap=False)
-    def name(self) -> str:
-        """"""
-
-        return self.joint_description.name
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def position(self, dof: int | None = None) -> jtp.Float:
-        """"""
-
-        dof = dof if dof is not None else 0
-
-        return jnp.array(
-            self.parent_model.joint_positions(joint_names=(self.name(),))[dof],
-            dtype=float,
-        )
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def velocity(self, dof: int | None = None) -> jtp.Float:
-        """"""
-
-        dof = dof if dof is not None else 0
-
-        return jnp.array(
-            self.parent_model.joint_velocities(joint_names=(self.name(),))[dof],
-            dtype=float,
-        )
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def force_target(self, dof: int | None = None) -> jtp.Float:
-        """"""
-
-        dof = dof if dof is not None else 0
-
-        return jnp.array(
-            self.parent_model.joint_generalized_forces_targets(
-                joint_names=(self.name(),)
-            )[dof],
-            dtype=float,
-        )
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def position_limit(self, dof: int | None = None) -> tuple[jtp.Float, jtp.Float]:
-        """"""
-
-        dof = dof if dof is not None else 0
-
-        if not_tracing(dof) and dof != 0:
-            msg = "Only joints with 1 DoF are currently supported"
-            raise ValueError(msg)
-
-        low, high = self.joint_description.position_limit
-
-        return jnp.array(low, dtype=float), jnp.array(high, dtype=float)
-
-    # =============
-    # Motor methods
-    # =============
-    @functools.partial(oop.jax_tf.method_ro)
-    def motor_inertia(self) -> jtp.Vector:
-        """"""
-
-        return jnp.array(self.joint_description.motor_inertia, dtype=float)
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def motor_gear_ratio(self) -> jtp.Vector:
-        """"""
-
-        return jnp.array(self.joint_description.motor_gear_ratio, dtype=float)
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def motor_viscous_friction(self) -> jtp.Vector:
-        """"""
-
-        return jnp.array(self.joint_description.motor_viscous_friction, dtype=float)
-
-    # =================
-    # Multi-DoF methods
-    # =================
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def joint_position(self) -> jtp.Vector:
-        """"""
-
-        return self.parent_model.joint_positions(joint_names=(self.name(),))
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def joint_velocity(self) -> jtp.Vector:
-        """"""
-
-        return self.parent_model.joint_velocities(joint_names=(self.name(),))
-
-    @functools.partial(oop.jax_tf.method_ro)
-    def joint_force_target(self) -> jtp.Vector:
-        """"""
-
-        return self.parent_model.joint_generalized_forces_targets(
-            joint_names=(self.name(),)
-        )