Update hash and comparison of JaxSimModelData

Author: diegoferigo

src/jaxsim/api/data.py

@@ -45,12 +45,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
                 hash(self.state),
-                hash(tuple(self.gravity.flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.gravity),
                 hash(self.soft_contacts_params),
-                hash(jnp.atleast_1d(self.time_ns).flatten().tolist()),
+                hash(tuple(self.time_ns.flatten().tolist())),
             )
         )
 

src/jaxsim/api/ode_data.py

@@ -283,12 +283,16 @@ class PhysicsModelState(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         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())),
+                HashedNumpyArray.hash_of_array(self.joint_positions),
+                HashedNumpyArray.hash_of_array(self.joint_velocities),
+                HashedNumpyArray.hash_of_array(self.base_position),
+                HashedNumpyArray.hash_of_array(self.base_quaternion),
+                HashedNumpyArray.hash_of_array(self.base_linear_velocity),
+                HashedNumpyArray.hash_of_array(self.base_angular_velocity),
             )
         )
 
@@ -613,9 +617,9 @@ class SoftContactsState(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
-        return hash(
-            tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
-        )
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return HashedNumpyArray.hash_of_array(self.tangential_deformation)
 
     def __eq__(self, other: SoftContactsState) -> bool:
 

src/jaxsim/rbda/soft_contacts.py

@@ -31,11 +31,13 @@ class SoftContactsParams(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         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())),
+                HashedNumpyArray.hash_of_array(self.K),
+                HashedNumpyArray.hash_of_array(self.D),
+                HashedNumpyArray.hash_of_array(self.mu),
             )
         )