jaxsim 0.2.1.dev80__tar.gz → 0.2.1.dev98__tar.gz

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev80
+Version: 0.2.1.dev98
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/pyproject.toml

@@ -24,13 +24,10 @@ profile = "black"
 [tool.pytest.ini_options]
 addopts = "-rsxX -v --strict-markers"
 minversion = "6.0"
-preview = true
 testpaths = [
   "tests",
 ]
 
-target-version = "py311"
-
 [tool.ruff]
 exclude = [
   ".git",

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev80'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev80')
+__version__ = version = '0.2.1.dev98'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev98')

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/api/data.py

@@ -30,7 +30,7 @@ except ImportError:
 @jax_dataclasses.pytree_dataclass
 class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     """
-    Class containing the state of a `JaxSimModel` object.
+    Class containing the data of a `JaxSimModel` object.
     """
 
     state: ODEState
@@ -43,6 +43,24 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
     )
 
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.state),
+                hash(tuple(self.gravity.flatten().tolist())),
+                hash(self.soft_contacts_params),
+                hash(jnp.atleast_1d(self.time_ns).flatten().tolist()),
+            )
+        )
+
+    def __eq__(self, other: JaxSimModelData) -> bool:
+
+        if not isinstance(other, JaxSimModelData):
+            return False
+
+        return hash(self) == hash(other)
+
     def valid(self, model: js.model.JaxSimModel | None = None) -> bool:
         """
         Check if the current state is valid for the given model.

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/api/frame.py

@@ -30,7 +30,7 @@ def idx_of_parent_link(model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike) -
     """
 
     # Get the intermediate representation parsed from the model description.
-    ir = model.description.get()
+    ir = model.description
 
     # Extract the indices of the frame and the link it is attached to.
     F = ir.frames[frame_idx - model.number_of_links()]
@@ -51,7 +51,7 @@ def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> int:
         The index of the frame.
     """
 
-    frame_names = np.array([frame.name for frame in model.description.get().frames])
+    frame_names = np.array([frame.name for frame in model.description.frames])
 
     if frame_name in frame_names:
         idx_in_list = np.argwhere(frame_names == frame_name)
@@ -72,7 +72,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str
         The name of the frame.
     """
 
-    return model.description.get().frames[frame_index - model.number_of_links()].name
+    return model.description.frames[frame_index - model.number_of_links()].name
 
 
 @functools.partial(jax.jit, static_argnames=["frame_names"])
@@ -144,7 +144,7 @@ def transform(
     W_H_L = js.link.transform(model=model, data=data, link_index=L)
 
     # Get the static frame pose wrt the parent link.
-    frame = model.description.get().frames[frame_index - model.number_of_links()]
+    frame = model.description.frames[frame_index - model.number_of_links()]
     L_H_F = frame.pose
 
     # Combine the transforms computing the frame pose.

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/api/kin_dyn_parameters.py

@@ -11,7 +11,7 @@ from jax_dataclasses import Static
 import jaxsim.typing as jtp
 from jaxsim.math import Inertia, JointModel, supported_joint_motion
 from jaxsim.parsers.descriptions import JointDescription, ModelDescription
-from jaxsim.utils import JaxsimDataclass
+from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
 
 @jax_dataclasses.pytree_dataclass
@@ -32,8 +32,8 @@ class KynDynParameters(JaxsimDataclass):
 
     # Static
     link_names: Static[tuple[str]]
-    parent_array: Static[jtp.Vector]
-    support_body_array_bool: Static[jtp.Matrix]
+    _parent_array: Static[HashedNumpyArray]
+    _support_body_array_bool: Static[HashedNumpyArray]
 
     # Links
     link_parameters: LinkParameters
@@ -45,6 +45,14 @@ class KynDynParameters(JaxsimDataclass):
     joint_model: JointModel
     joint_parameters: JointParameters | None
 
+    @property
+    def parent_array(self) -> jtp.Vector:
+        return self._parent_array.get()
+
+    @property
+    def support_body_array_bool(self) -> jtp.Matrix:
+        return self._support_body_array_bool.get()
+
     @staticmethod
     def build(model_description: ModelDescription) -> KynDynParameters:
         """
@@ -191,8 +199,8 @@ class KynDynParameters(JaxsimDataclass):
 
         return KynDynParameters(
             link_names=tuple(l.name for l in ordered_links),
-            parent_array=parent_array,
-            support_body_array_bool=support_body_array_bool,
+            _parent_array=HashedNumpyArray(array=parent_array),
+            _support_body_array_bool=HashedNumpyArray(array=support_body_array_bool),
             link_parameters=link_parameters,
             joint_model=joint_model,
             joint_parameters=joint_parameters,
@@ -204,23 +212,18 @@ class KynDynParameters(JaxsimDataclass):
         if not isinstance(other, KynDynParameters):
             return False
 
-        equal = True
-        equal = equal and self.number_of_links() == other.number_of_links()
-        equal = equal and self.number_of_joints() == other.number_of_joints()
-        equal = equal and jnp.allclose(self.parent_array, other.parent_array)
-
-        return equal
+        return hash(self) == hash(other)
 
     def __hash__(self) -> int:
 
-        h = (
-            hash(self.number_of_links()),
-            hash(self.number_of_joints()),
-            hash(tuple(self.parent_array.tolist())),
+        return hash(
+            (
+                hash(self.number_of_links()),
+                hash(self.number_of_joints()),
+                hash(tuple(jnp.atleast_1d(self.parent_array).flatten().tolist())),
+            )
         )
 
-        return hash(h)
-
     # =============================
     # Helpers to extract parameters
     # =============================
@@ -388,7 +391,7 @@ class KynDynParameters(JaxsimDataclass):
             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(self.joint_model.joint_axis),
+                jnp.array([j.axis for j in self.joint_model.joint_axis]),
             )
 
         # Extract the transforms and motion subspaces of the joints.

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/api/model.py

@@ -32,18 +32,37 @@ class JaxSimModel(JaxsimDataclass):
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
         default=jaxsim.terrain.FlatTerrain(), repr=False, compare=False, hash=False
     )
+    kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
+        dataclasses.field(default=None, repr=False, compare=False, hash=False)
+    )
 
     built_from: Static[str | pathlib.Path | rod.Model | None] = dataclasses.field(
         default=None, repr=False, compare=False, hash=False
     )
 
-    description: Static[
+    _description: Static[
         HashlessObject[jaxsim.parsers.descriptions.ModelDescription | None]
     ] = dataclasses.field(default=None, repr=False, compare=False, hash=False)
 
-    kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
-        dataclasses.field(default=None, repr=False, compare=False, hash=False)
-    )
+    @property
+    def description(self) -> jaxsim.parsers.descriptions.ModelDescription:
+        return self._description.get()
+
+    def __eq__(self, other: JaxSimModel) -> bool:
+
+        if not isinstance(other, JaxSimModel):
+            return False
+
+        return hash(self) == hash(other)
+
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.model_name),
+                hash(self.kin_dyn_parameters),
+            )
+        )
 
     # ========================
     # Initialization and state
@@ -137,7 +156,7 @@ class JaxSimModel(JaxsimDataclass):
         # Build the model
         model = JaxSimModel(
             model_name=model_name,
-            description=HashlessObject(obj=model_description),
+            _description=HashlessObject(obj=model_description),
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
@@ -260,7 +279,7 @@ class JaxSimModel(JaxsimDataclass):
             The names of the links in the model.
         """
 
-        return tuple([frame.name for frame in self.description.get().frames])
+        return tuple(frame.name for frame in self.description.frames)
 
 
 # =====================
@@ -297,7 +316,7 @@ def reduce(
 
     # Copy the model description with a deep copy of the joints.
     intermediate_description = dataclasses.replace(
-        model.description.get(), joints=copy.deepcopy(model.description.get().joints)
+        model.description, joints=copy.deepcopy(model.description.joints)
     )
 
     # Update the initial position of the joints.

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/api/ode_data.py

@@ -281,6 +281,24 @@ class PhysicsModelState(JaxsimDataclass):
         default_factory=lambda: jnp.zeros(3)
     )
 
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(tuple(jnp.atleast_1d(self.joint_positions.flatten().tolist()))),
+                hash(tuple(jnp.atleast_1d(self.joint_velocities.flatten().tolist()))),
+                hash(tuple(self.base_position.flatten().tolist())),
+                hash(tuple(self.base_quaternion.flatten().tolist())),
+            )
+        )
+
+    def __eq__(self, other: PhysicsModelState) -> bool:
+
+        if not isinstance(other, PhysicsModelState):
+            return False
+
+        return hash(self) == hash(other)
+
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
@@ -593,6 +611,19 @@ class SoftContactsState(JaxsimDataclass):
 
     tangential_deformation: jtp.Matrix
 
+    def __hash__(self) -> int:
+
+        return hash(
+            tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
+        )
+
+    def __eq__(self, other: SoftContactsState) -> bool:
+
+        if not isinstance(other, SoftContactsState):
+            return False
+
+        return hash(self) == hash(other)
+
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/math/joint_model.py

@@ -39,7 +39,7 @@ class JointModel:
 
     joint_dofs: Static[tuple[int, ...]]
     joint_names: Static[tuple[str, ...]]
-    joint_types: Static[tuple[JointType, ...]]
+    joint_types: Static[tuple[int, ...]]
     joint_axis: Static[tuple[JointGenericAxis, ...]]
 
     @staticmethod
@@ -109,7 +109,7 @@ class JointModel:
             joint_dofs=tuple([base_dofs] + [int(1) for _ in ordered_joints]),
             joint_names=tuple(["world_to_base"] + [j.name for j in ordered_joints]),
             joint_types=tuple([JointType.Fixed] + [j.jtype for j in ordered_joints]),
-            joint_axis=tuple([j.axis for j in ordered_joints]),
+            joint_axis=tuple(JointGenericAxis(axis=j.axis) for j in ordered_joints),
         )
 
     def parent_H_child(
@@ -201,7 +201,7 @@ class JointModel:
         pre_H_suc, S = supported_joint_motion(
             self.joint_types[joint_index],
             joint_position,
-            self.joint_axis[joint_index],
+            self.joint_axis[joint_index].axis,
         )
 
         return pre_H_suc, S
@@ -224,9 +224,9 @@ class JointModel:
 
 @jax.jit
 def supported_joint_motion(
-    joint_type: JointType,
+    joint_type: jtp.IntLike,
     joint_position: jtp.VectorLike,
-    joint_axis: JointGenericAxis,
+    joint_axis: jtp.VectorLike | None = None,
     /,
 ) -> tuple[jtp.Matrix, jtp.Array]:
     """
@@ -234,8 +234,8 @@ def supported_joint_motion(
 
     Args:
         joint_type: The type of the joint.
-        joint_axis: The axis of rotation or translation of the joint.
         joint_position: The position of the joint.
+        joint_axis: The optional 3D axis of rotation or translation of the joint.
 
     Returns:
         A tuple containing the homogeneous transformation and the motion subspace.
@@ -244,26 +244,33 @@ def supported_joint_motion(
     # Prepare the joint position
     s = jnp.array(joint_position).astype(float)
 
-    def compute_F():
+    def compute_F() -> tuple[jtp.Matrix, jtp.Array]:
         return jaxlie.SE3.identity(), jnp.zeros(shape=(6, 1))
 
-    def compute_R():
+    def compute_R() -> tuple[jtp.Matrix, jtp.Array]:
+
+        # Get the additional argument specifying the joint axis.
+        # This is a metadata required by only some joint types.
+        axis = jnp.array(joint_axis).astype(float).squeeze()
+
         pre_H_suc = jaxlie.SE3.from_rotation(
-            rotation=jaxlie.SO3.from_matrix(
-                Rotation.from_axis_angle(vector=s * joint_axis)
-            )
+            rotation=jaxlie.SO3.from_matrix(Rotation.from_axis_angle(vector=s * axis))
         )
 
-        S = jnp.vstack(jnp.hstack([jnp.zeros(3), joint_axis.squeeze()]))
+        S = jnp.vstack(jnp.hstack([jnp.zeros(3), axis]))
+
         return pre_H_suc, S
 
-    def compute_P():
-        pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
-            rotation=jaxlie.SO3.identity(),
-            translation=jnp.array(s * joint_axis),
-        )
+    def compute_P() -> tuple[jtp.Matrix, jtp.Array]:
+
+        # Get the additional argument specifying the joint axis.
+        # This is a metadata required by only some joint types.
+        axis = jnp.array(joint_axis).astype(float).squeeze()
+
+        pre_H_suc = jaxlie.SE3.from_translation(translation=jnp.array(s * axis))
+
+        S = jnp.vstack(jnp.hstack([axis, jnp.zeros(3)]))
 
-        S = jnp.vstack(jnp.hstack([joint_axis.squeeze(), jnp.zeros(3)]))
         return pre_H_suc, S
 
     pre_H_suc, S = jax.lax.switch(

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/parsers/descriptions/joint.py

@@ -30,9 +30,11 @@ class JointGenericAxis:
     axis: jtp.Vector
 
     def __hash__(self) -> int:
-        return hash((tuple(np.array(self.axis).tolist())))
+
+        return hash(tuple(self.axis.tolist()))
 
     def __eq__(self, other: JointGenericAxis) -> bool:
+
         if not isinstance(other, JointGenericAxis):
             return False
 

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/rbda/soft_contacts.py

@@ -29,6 +29,23 @@ class SoftContactsParams(JaxsimDataclass):
         default_factory=lambda: jnp.array(0.5, dtype=float)
     )
 
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(tuple(jnp.atleast_1d(self.K).flatten().tolist())),
+                hash(tuple(jnp.atleast_1d(self.D).flatten().tolist())),
+                hash(tuple(jnp.atleast_1d(self.mu).flatten().tolist())),
+            )
+        )
+
+    def __eq__(self, other: SoftContactsParams) -> bool:
+
+        if not isinstance(other, SoftContactsParams):
+            return NotImplemented
+
+        return hash(self) == hash(other)
+
     @staticmethod
     def build(
         K: jtp.FloatLike = 1e6, D: jtp.FloatLike = 2_000, mu: jtp.FloatLike = 0.5

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim/utils/__init__.py

@@ -1,5 +1,5 @@
 from jax_dataclasses._copy_and_mutate import _Mutability as Mutability
 
-from .hashless import HashlessObject
 from .jaxsim_dataclass import JaxsimDataclass
 from .tracing import not_tracing, tracing
+from .wrappers import HashedNumpyArray, HashlessObject

jaxsim-0.2.1.dev98/src/jaxsim/utils/wrappers.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import Generic, TypeVar
+
+import jax
+import jax_dataclasses
+import numpy as np
+import numpy.typing as npt
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class HashlessObject(Generic[T]):
+    """
+    A class that wraps an object and makes it hashless.
+
+    This is useful for creating particular JAX pytrees.
+    For example, to create a pytree with a static leaf that is ignored
+    by JAX when it compares two instances to trigger a JIT recompilation.
+    """
+
+    obj: T
+
+    def get(self: HashlessObject[T]) -> T:
+        return self.obj
+
+    def __hash__(self) -> int:
+
+        return 0
+
+    def __eq__(self, other: HashlessObject[T]) -> bool:
+
+        if not isinstance(other, HashlessObject) and isinstance(
+            other.get(), type(self.get())
+        ):
+            return False
+
+        return hash(self) == hash(other)
+
+
+@jax_dataclasses.pytree_dataclass
+class HashedNumpyArray:
+    """
+    A class that wraps a numpy array and makes it hashable.
+
+    This is useful for creating particular JAX pytrees.
+    For example, to create a pytree with a plain NumPy or JAX NumPy array as static leaf.
+
+    Note:
+        Calculating with the wrapper class the hash of a very large array can be
+        very expensive. If the array is large and only the equality operator is needed,
+        set `large_array=True` to use a faster comparison method.
+    """
+
+    array: jax.Array | npt.NDArray
+
+    large_array: jax_dataclasses.Static[bool] = dataclasses.field(
+        default=False, repr=False, compare=False, hash=False
+    )
+
+    def get(self) -> jax.Array | npt.NDArray:
+        return self.array
+
+    def __hash__(self) -> int:
+
+        return hash(tuple(np.atleast_1d(self.array).flatten().tolist()))
+
+    def __eq__(self, other: HashedNumpyArray) -> bool:
+
+        if not isinstance(other, HashedNumpyArray):
+            return False
+
+        if self.large_array:
+            return np.array_equal(self.array, other.array)
+
+        return hash(self) == hash(other)

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev80
+Version: 0.2.1.dev98
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/src/jaxsim.egg-info/SOURCES.txt

@@ -99,9 +99,9 @@ src/jaxsim/rbda/utils.py
 src/jaxsim/terrain/__init__.py
 src/jaxsim/terrain/terrain.py
 src/jaxsim/utils/__init__.py
-src/jaxsim/utils/hashless.py
 src/jaxsim/utils/jaxsim_dataclass.py
 src/jaxsim/utils/tracing.py
+src/jaxsim/utils/wrappers.py
 tests/__init__.py
 tests/conftest.py
 tests/test_api_com.py

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/tests/conftest.py

@@ -177,18 +177,19 @@ def jaxsim_model_sphere() -> js.model.JaxSimModel:
 
 
 @pytest.fixture(scope="session")
-def jaxsim_model_ergocub() -> js.model.JaxSimModel:
+def ergocub_model_description_path() -> pathlib.Path:
     """
-    Fixture providing the JaxSim model of the ErgoCub robot.
+    Fixture providing the path to the URDF model description of the ErgoCub robot.
 
     Returns:
-        The JaxSim model of the ErgoCub robot.
+        The path to the URDF model description of the ErgoCub robot.
     """
 
     try:
         os.environ["ROBOT_DESCRIPTION_COMMIT"] = "v0.7.1"
 
         import robot_descriptions.ergocub_description
+
     finally:
         _ = os.environ.pop("ROBOT_DESCRIPTION_COMMIT", None)
 
@@ -198,7 +199,21 @@ def jaxsim_model_ergocub() -> js.model.JaxSimModel:
         )
     )
 
-    return build_jaxsim_model(model_description=model_urdf_path)
+    return model_urdf_path
+
+
+@pytest.fixture(scope="session")
+def jaxsim_model_ergocub(
+    ergocub_model_description_path: pathlib.Path,
+) -> js.model.JaxSimModel:
+    """
+    Fixture providing the JaxSim model of the ErgoCub robot.
+
+    Returns:
+        The JaxSim model of the ErgoCub robot.
+    """
+
+    return build_jaxsim_model(model_description=ergocub_model_description_path)
 
 
 @pytest.fixture(scope="session")

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/tests/test_api_com.py

@@ -15,7 +15,7 @@ def test_com_properties(
 
     model = jaxsim_models_types
 
-    key, subkey = jax.random.split(prng_key, num=2)
+    _, subkey = jax.random.split(prng_key, num=2)
     data = js.data.random_model_data(
         model=model, key=subkey, velocity_representation=velocity_representation
     )

{jaxsim-0.2.1.dev80 → jaxsim-0.2.1.dev98}/tests/test_api_data.py

@@ -27,7 +27,7 @@ def test_data_joint_indexing(
 
     model = jaxsim_models_types
 
-    key, subkey = jax.random.split(prng_key, num=2)
+    _, subkey = jax.random.split(prng_key, num=2)
     data = js.data.random_model_data(
         model=model, key=subkey, velocity_representation=velocity_representation
     )
@@ -56,7 +56,7 @@ def test_data_switch_velocity_representation(
 
     model = jaxsim_models_types
 
-    key, subkey = jax.random.split(prng_key, num=2)
+    _, subkey = jax.random.split(prng_key, num=2)
     data = js.data.random_model_data(
         model=model, key=subkey, velocity_representation=VelRepr.Inertial
     )
@@ -98,7 +98,7 @@ def test_data_change_velocity_representation(
 
     model = jaxsim_models_types
 
-    key, subkey = jax.random.split(prng_key, num=2)
+    _, subkey = jax.random.split(prng_key, num=2)
     data = js.data.random_model_data(
         model=model, key=subkey, velocity_representation=VelRepr.Inertial
     )