jaxsim 0.5.1.dev126__py3-none-any.whl → 0.5.1.dev139__py3-none-any.whl

jaxsim/mujoco/model.py

@@ -254,17 +254,17 @@ class MujocoModelHelper:
     # ==================
 
     def number_of_joints(self) -> int:
-        """Returns the number of joints in the model."""
+        """Return the number of joints in the model."""
 
         return self.model.njnt
 
     def number_of_dofs(self) -> int:
-        """Returns the number of DoFs in the model."""
+        """Return the number of DoFs in the model."""
 
         return self.model.nq
 
     def joint_names(self) -> list[str]:
-        """Returns the names of the joints in the model."""
+        """Return the names of the joints in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_JOINT, idx)
@@ -272,7 +272,7 @@ class MujocoModelHelper:
         ]
 
     def joint_dofs(self, joint_name: str) -> int:
-        """Returns the number of DoFs of a joint."""
+        """Return the number of DoFs of a joint."""
 
         if joint_name not in self.joint_names():
             raise ValueError(f"Joint '{joint_name}' not found")
@@ -280,7 +280,7 @@ class MujocoModelHelper:
         return self.data.joint(joint_name).qpos.size
 
     def joint_position(self, joint_name: str) -> npt.NDArray:
-        """Returns the position of a joint."""
+        """Return the position of a joint."""
 
         if joint_name not in self.joint_names():
             raise ValueError(f"Joint '{joint_name}' not found")
@@ -288,7 +288,7 @@ class MujocoModelHelper:
         return self.data.joint(joint_name).qpos
 
     def joint_positions(self, joint_names: list[str] | None = None) -> npt.NDArray:
-        """Returns the positions of the joints."""
+        """Return the positions of the joints."""
 
         joint_names = joint_names if joint_names is not None else self.joint_names()
 
@@ -299,7 +299,7 @@ class MujocoModelHelper:
     def set_joint_position(
         self, joint_name: str, position: npt.NDArray | float
     ) -> None:
-        """Sets the position of a joint."""
+        """Set the position of a joint."""
 
         position = np.atleast_1d(np.array(position).squeeze())
 
@@ -328,12 +328,12 @@ class MujocoModelHelper:
     # ==================
 
     def number_of_bodies(self) -> int:
-        """Returns the number of bodies in the model."""
+        """Return the number of bodies in the model."""
 
         return self.model.nbody
 
     def body_names(self) -> list[str]:
-        """Returns the names of the bodies in the model."""
+        """Return the names of the bodies in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_BODY, idx)
@@ -341,7 +341,7 @@ class MujocoModelHelper:
         ]
 
     def body_position(self, body_name: str) -> npt.NDArray:
-        """Returns the position of a body."""
+        """Return the position of a body."""
 
         if body_name not in self.body_names():
             raise ValueError(f"Body '{body_name}' not found")
@@ -349,7 +349,7 @@ class MujocoModelHelper:
         return self.data.body(body_name).xpos
 
     def body_orientation(self, body_name: str, dcm: bool = False) -> npt.NDArray:
-        """Returns the orientation of a body."""
+        """Return the orientation of a body."""
 
         if body_name not in self.body_names():
             raise ValueError(f"Body '{body_name}' not found")
@@ -363,12 +363,12 @@ class MujocoModelHelper:
     # ======================
 
     def number_of_geometries(self) -> int:
-        """Returns the number of geometries in the model."""
+        """Return the number of geometries in the model."""
 
         return self.model.ngeom
 
     def geometry_names(self) -> list[str]:
-        """Returns the names of the geometries in the model."""
+        """Return the names of the geometries in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_GEOM, idx)
@@ -376,7 +376,7 @@ class MujocoModelHelper:
         ]
 
     def geometry_position(self, geometry_name: str) -> npt.NDArray:
-        """Returns the position of a geometry."""
+        """Return the position of a geometry."""
 
         if geometry_name not in self.geometry_names():
             raise ValueError(f"Geometry '{geometry_name}' not found")
@@ -386,7 +386,7 @@ class MujocoModelHelper:
     def geometry_orientation(
         self, geometry_name: str, dcm: bool = False
     ) -> npt.NDArray:
-        """Returns the orientation of a geometry."""
+        """Return the orientation of a geometry."""
 
         if geometry_name not in self.geometry_names():
             raise ValueError(f"Geometry '{geometry_name}' not found")

jaxsim/mujoco/utils.py

@@ -133,6 +133,9 @@ class MujocoCamera:
 
     @classmethod
     def build(cls, **kwargs) -> MujocoCamera:
+        """
+        Build a Mujoco camera from a dictionary.
+        """
 
         if not all(isinstance(value, str) for value in kwargs.values()):
             raise ValueError(f"Values must be strings: {kwargs}")
@@ -219,5 +222,7 @@ class MujocoCamera:
         )
 
     def asdict(self) -> dict[str, str]:
-
+        """
+        Convert the camera to a dictionary.
+        """
         return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}

jaxsim/mujoco/visualizer.py

@@ -10,7 +10,9 @@ import numpy.typing as npt
 
 
 class MujocoVideoRecorder:
-    """"""
+    """
+    Video recorder for the MuJoCo passive viewer.
+    """
 
     def __init__(
         self,
@@ -64,7 +66,7 @@ class MujocoVideoRecorder:
         self.model = model if model is not None else self.model
 
     def render_frame(self, camera_name: str = "track") -> npt.NDArray:
-        """Renders a frame."""
+        """Render a frame."""
 
         mujoco.mj_forward(self.model, self.data)
         self.renderer.update_scene(data=self.data, camera=camera_name)
@@ -72,13 +74,13 @@ class MujocoVideoRecorder:
         return self.renderer.render()
 
     def record_frame(self, camera_name: str = "track") -> None:
-        """Stores a frame in the buffer."""
+        """Store a frame in the buffer."""
 
         frame = self.render_frame(camera_name=camera_name)
         self.frames.append(frame)
 
     def write_video(self, path: pathlib.Path, exist_ok: bool = False) -> None:
-        """Writes the video to a file."""
+        """Write the video to a file."""
 
         # Resolve the path to the video.
         path = path.expanduser().resolve()
@@ -117,7 +119,9 @@ class MujocoVideoRecorder:
 
 
 class MujocoVisualizer:
-    """"""
+    """
+    Visualizer for the MuJoCo passive viewer.
+    """
 
     def __init__(
         self, model: mj.MjModel | None = None, data: mj.MjData | None = None
@@ -139,7 +143,7 @@ class MujocoVisualizer:
         model: mj.MjModel | None = None,
         data: mj.MjData | None = None,
     ) -> None:
-        """Updates the viewer with the current model and data."""
+        """Update the viewer with the current model and data."""
 
         data = data if data is not None else self.data
         model = model if model is not None else self.model
@@ -150,7 +154,7 @@ class MujocoVisualizer:
     def open_viewer(
         self, model: mj.MjModel | None = None, data: mj.MjData | None = None
     ) -> mj.viewer.Handle:
-        """Opens a viewer."""
+        """Open a viewer."""
 
         data = data if data is not None else self.data
         model = model if model is not None else self.model

jaxsim/parsers/descriptions/collision.py

@@ -22,7 +22,6 @@ class CollidablePoint:
         parent_link: The parent link to which the collidable point is attached.
         position: The position of the collidable point relative to the parent link.
         enabled: A flag indicating whether the collidable point is enabled for collision detection.
-
     """
 
     parent_link: LinkDescription
@@ -86,7 +85,6 @@ class CollisionShape(abc.ABC):
 
     Attributes:
         collidable_points: A list of collidable points associated with the collision shape.
-
     """
 
     collidable_points: tuple[CollidablePoint]
@@ -107,7 +105,6 @@ class BoxCollision(CollisionShape):
 
     Attributes:
         center: The center of the box in the local frame of the collision shape.
-
     """
 
     center: jtp.VectorLike
@@ -135,7 +132,6 @@ class SphereCollision(CollisionShape):
 
     Attributes:
         center: The center of the sphere in the local frame of the collision shape.
-
     """
 
     center: jtp.VectorLike
@@ -158,6 +154,13 @@ class SphereCollision(CollisionShape):
 
 @dataclasses.dataclass
 class MeshCollision(CollisionShape):
+    """
+    Represents a mesh-shaped collision shape.
+
+    Attributes:
+        center: The center of the mesh in the local frame of the collision shape.
+    """
+
     center: jtp.VectorLike
 
     def __hash__(self) -> int:

jaxsim/parsers/descriptions/joint.py

@@ -14,6 +14,9 @@ from .link import LinkDescription
 
 @dataclasses.dataclass(frozen=True)
 class JointType:
+    """
+    Enumeration of joint types.
+    """
 
     Fixed: ClassVar[int] = 0
     Revolute: ClassVar[int] = 1
@@ -47,20 +50,19 @@ class JointDescription(JaxsimDataclass):
     In-memory description of a robot link.
 
     Attributes:
-        name (str): The name of the joint.
-        axis (npt.NDArray): The axis of rotation or translation for the joint.
-        pose (npt.NDArray): The pose transformation matrix of the joint.
-        jtype (JointType): The type of the joint.
-        child (LinkDescription): The child link attached to the joint.
-        parent (LinkDescription): The parent link attached to the joint.
-        index (Optional[int]): An optional index for the joint.
-        friction_static (float): The static friction coefficient for the joint.
-        friction_viscous (float): The viscous friction coefficient for the joint.
-        position_limit_damper (float): The damper coefficient for position limits.
-        position_limit_spring (float): The spring coefficient for position limits.
-        position_limit (Tuple[float, float]): The position limits for the joint.
-        initial_position (Union[float, npt.NDArray]): The initial position of the joint.
-
+        name: The name of the joint.
+        axis: The axis of rotation or translation for the joint.
+        pose: The pose transformation matrix of the joint.
+        jtype: The type of the joint.
+        child: The child link attached to the joint.
+        parent: The parent link attached to the joint.
+        index: An optional index for the joint.
+        friction_static: The static friction coefficient for the joint.
+        friction_viscous: The viscous friction coefficient for the joint.
+        position_limit_damper: The damper coefficient for position limits.
+        position_limit_spring: The spring coefficient for position limits.
+        position_limit: The position limits for the joint.
+        initial_position: The initial position of the joint.
     """
 
     name: jax_dataclasses.Static[str]

jaxsim/parsers/descriptions/model.py

@@ -158,7 +158,7 @@ class ModelDescription(KinematicGraph):
         Reduce the model by removing specified joints.
 
         Args:
-            The joint names to consider.
+            considered_joints: Sequence of joint names to consider.
 
         Returns:
             A `ModelDescription` instance that only includes the considered joints.

jaxsim/parsers/kinematic_graph.py

@@ -97,20 +97,32 @@ class KinematicGraph(Sequence[LinkDescription]):
 
     @functools.cached_property
     def links_dict(self) -> dict[str, LinkDescription]:
+        """
+        Get a dictionary of links indexed by their name.
+        """
         return {l.name: l for l in iter(self)}
 
     @functools.cached_property
     def frames_dict(self) -> dict[str, LinkDescription]:
+        """
+        Get a dictionary of frames indexed by their name.
+        """
         return {f.name: f for f in self.frames}
 
     @functools.cached_property
     def joints_dict(self) -> dict[str, JointDescription]:
+        """
+        Get a dictionary of joints indexed by their name.
+        """
         return {j.name: j for j in self.joints}
 
     @functools.cached_property
     def joints_connection_dict(
         self,
     ) -> dict[tuple[str, str], JointDescription]:
+        """
+        Get a dictionary of joints indexed by the tuple (parent, child) link names.
+        """
         return {(j.parent.name, j.child.name): j for j in self.joints}
 
     def __post_init__(self) -> None:
@@ -734,9 +746,15 @@ class KinematicGraph(Sequence[LinkDescription]):
         raise TypeError(type(key).__name__)
 
     def count(self, value: LinkDescription) -> int:
+        """
+        Count the occurrences of a link in the kinematic graph.
+        """
         return list(iter(self)).count(value)
 
     def index(self, value: LinkDescription, start: int = 0, stop: int = -1) -> int:
+        """
+        Find the index of a link in the kinematic graph.
+        """
         return list(iter(self)).index(value, start, stop)
 
 
@@ -747,6 +765,12 @@ class KinematicGraph(Sequence[LinkDescription]):
 
 @dataclasses.dataclass(frozen=True)
 class KinematicGraphTransforms:
+    """
+    Class to compute forward kinematics on a kinematic graph.
+
+    Attributes:
+        graph: The kinematic graph on which to compute forward kinematics.
+    """
 
     graph: KinematicGraph
 
@@ -767,6 +791,9 @@ class KinematicGraphTransforms:
 
     @property
     def initial_joint_positions(self) -> npt.NDArray:
+        """
+        Get the initial joint positions of the kinematic graph.
+        """
 
         return np.atleast_1d(
             np.array(list(self._initial_joint_positions.values()))
@@ -910,6 +937,17 @@ class KinematicGraphTransforms:
         joint_axis: npt.NDArray,
         joint_position: float | None = None,
     ) -> npt.NDArray:
+        """
+        Compute the SE(3) transform from the predecessor to the successor frame.
+
+        Args:
+            joint_type: The type of the joint.
+            joint_axis: The axis of the joint.
+            joint_position: The position of the joint.
+
+        Returns:
+            The 4x4 transform matrix from the predecessor to the successor frame.
+        """
 
         import jaxsim.math
 

jaxsim/parsers/rod/meshes.py

@@ -6,14 +6,14 @@ VALID_AXIS = {"x": 0, "y": 1, "z": 2}
 
 def extract_points_vertices(mesh: trimesh.Trimesh) -> np.ndarray:
     """
-    Extracts the vertices of a mesh as points.
+    Extract the vertices of a mesh as points.
     """
     return mesh.vertices
 
 
 def extract_points_random_surface_sampling(mesh: trimesh.Trimesh, n) -> np.ndarray:
     """
-    Extracts N random points from the surface of a mesh.
+    Extract N random points from the surface of a mesh.
 
     Args:
         mesh: The mesh from which to extract points.
@@ -30,7 +30,7 @@ def extract_points_uniform_surface_sampling(
     mesh: trimesh.Trimesh, n: int
 ) -> np.ndarray:
     """
-    Extracts N uniformly sampled points from the surface of a mesh.
+    Extract N uniformly sampled points from the surface of a mesh.
 
     Args:
         mesh: The mesh from which to extract points.
@@ -47,7 +47,7 @@ def extract_points_select_points_over_axis(
     mesh: trimesh.Trimesh, axis: str, direction: str, n: int
 ) -> np.ndarray:
     """
-    Extracts N points from a mesh along a specified axis. The points are selected based on their position along the axis.
+    Extract N points from a mesh along a specified axis. The points are selected based on their position along the axis.
 
     Args:
         mesh: The mesh from which to extract points.
@@ -75,7 +75,7 @@ def extract_points_aap(
     lower: float | None = None,
 ) -> np.ndarray:
     """
-    Extracts points from a mesh along a specified axis within a specified range. The points are selected based on their position along the axis.
+    Extract points from a mesh along a specified axis within a specified range. The points are selected based on their position along the axis.
 
     Args:
         mesh: The mesh from which to extract points.

jaxsim/parsers/rod/parser.py

@@ -364,7 +364,7 @@ def build_model_description(
     is_urdf: bool | None = None,
 ) -> descriptions.ModelDescription:
     """
-    Builds a model description from an SDF/URDF resource.
+    Build a model description from an SDF/URDF resource.
 
     Args:
         model_description: A path to an SDF/URDF file, a string containing its content,

jaxsim/parsers/rod/utils.py

@@ -223,6 +223,17 @@ def create_mesh_collision(
     link_description: descriptions.LinkDescription,
     method: MeshMappingMethod = None,
 ) -> descriptions.MeshCollision:
+    """
+    Create a mesh collision from an SDF collision element.
+
+    Args:
+        collision: The SDF collision element.
+        link_description: The link description.
+        method: The method to use for mesh wrapping.
+
+    Returns:
+        The mesh collision description.
+    """
 
     file = pathlib.Path(resolve_local_uri(uri=collision.geometry.mesh.uri))
     file_type = file.suffix.replace(".", "")

jaxsim/rbda/contacts/common.py

@@ -125,6 +125,7 @@ class ContactModel(JaxsimDataclass):
         Args:
             model: The robot model considered by the contact model.
             data: The data of the considered model.
+            **kwargs: Optional additional arguments, specific to the contact model.
 
         Returns:
             A tuple containing as first element the computed 6D contact force applied to
@@ -146,6 +147,7 @@ class ContactModel(JaxsimDataclass):
         Args:
             model: The robot model considered by the contact model.
             data: The data of the considered model.
+            **kwargs: Optional additional arguments, specific to the contact model.
 
         Returns:
             A tuple containing as first element the 6D contact force applied to the

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -112,7 +112,7 @@ class RelaxedRigidContactsParams(common.ContactsParams):
         damping: jtp.FloatLike | None = None,
         mu: jtp.FloatLike | None = None,
     ) -> Self:
-        """Create a `RelaxedRigidContactsParams` instance"""
+        """Create a `RelaxedRigidContactsParams` instance."""
 
         def default(name: str):
             return cls.__dataclass_fields__[name].default_factory()
@@ -160,6 +160,7 @@ class RelaxedRigidContactsParams(common.ContactsParams):
         )
 
     def valid(self) -> jtp.BoolLike:
+        """Check if the parameters are valid."""
 
         return bool(
             jnp.all(self.time_constant >= 0.0)
@@ -187,6 +188,7 @@ class RelaxedRigidContacts(common.ContactModel):
 
     @property
     def solver_options(self) -> dict[str, Any]:
+        """Get the solver options."""
 
         return dict(
             zip(
@@ -207,6 +209,7 @@ class RelaxedRigidContacts(common.ContactModel):
 
         Args:
             solver_options: The options to pass to the L-BFGS solver.
+            **kwargs: The parameters of the relaxed rigid contacts model.
 
         Returns:
             The `RelaxedRigidContacts` instance.
@@ -483,8 +486,8 @@ class RelaxedRigidContacts(common.ContactModel):
 
         Args:
             model: The jaxsim model.
-            penetration: The point position in the constraint frame.
-            velocity: The point velocity in the constraint frame.
+            position_constraint: The position of the collidable points in the constraint frame.
+            velocity_constraint: The velocity of the collidable points in the constraint frame.
             parameters: The parameters of the relaxed rigid contacts model.
 
         Returns:
@@ -526,7 +529,7 @@ class RelaxedRigidContacts(common.ContactModel):
             vel: jtp.Vector,
         ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector]:
             """
-            Calculates impedance and offset acceleration in constraint frame.
+            Calculate impedance and offset acceleration in constraint frame.
 
             Args:
                 pos: position in constraint frame.

jaxsim/rbda/contacts/rigid.py

@@ -62,7 +62,7 @@ class RigidContactsParams(ContactsParams):
         K: jtp.FloatLike | None = None,
         D: jtp.FloatLike | None = None,
     ) -> Self:
-        """Create a `RigidContactParams` instance"""
+        """Create a `RigidContactParams` instance."""
 
         return cls(
             mu=jnp.array(
@@ -79,7 +79,7 @@ class RigidContactsParams(ContactsParams):
         )
 
     def valid(self) -> jtp.BoolLike:
-
+        """Check if the parameters are valid."""
         return bool(
             jnp.all(self.mu >= 0.0)
             and jnp.all(self.K >= 0.0)
@@ -104,6 +104,7 @@ class RigidContacts(ContactModel):
 
     @property
     def solver_options(self) -> dict[str, Any]:
+        """Get the solver options as a dictionary."""
 
         return dict(
             zip(
@@ -127,6 +128,7 @@ class RigidContacts(ContactModel):
             regularization_delassus:
                 The regularization term to add to the diagonal of the Delassus matrix.
             solver_options: The options to pass to the QP solver.
+            **kwargs: Extra arguments which are ignored.
 
         Returns:
             The `RigidContacts` instance.
@@ -173,7 +175,8 @@ class RigidContacts(ContactModel):
         J_WC: jtp.MatrixLike,
         data: js.data.JaxSimModelData,
     ) -> jtp.Vector:
-        """Returns the new velocity of the system after a potential impact.
+        """
+        Return the new velocity of the system after a potential impact.
 
         Args:
             inactive_collidable_points: The activation state of the collidable points.
@@ -413,7 +416,8 @@ class RigidContacts(ContactModel):
         inactive_collidable_points: jtp.Vector, mu: jtp.FloatLike
     ) -> jtp.Matrix:
         """
-        Compute the inequality constraint matrix for a single collidable point
+        Compute the inequality constraint matrix for a single collidable point.
+
         Rows 0-3: enforce the friction pyramid constraint,
         Row 4: last one is for the non negativity of the vertical force
         Row 5: contact complementarity condition

jaxsim/rbda/contacts/soft.py

@@ -207,6 +207,7 @@ class SoftContacts(common.ContactModel):
             model:
                 The robot model considered by the contact model.
                 If passed, it is used to estimate good default parameters.
+            **kwargs: Additional parameters to pass to the contact model.
 
         Returns:
             The `SoftContacts` instance.
@@ -244,6 +245,28 @@ class SoftContacts(common.ContactModel):
         p: jtp.FloatLike = 0.5,
         q: jtp.FloatLike = 0.5,
     ) -> tuple[jtp.Vector, jtp.Vector]:
+        """
+        Compute the contact force using the Hunt/Crossley model.
+
+        Args:
+            position: The position of the collidable point.
+            velocity: The velocity of the collidable point.
+            tangential_deformation: The material deformation of the collidable point.
+            terrain: The terrain model.
+            K: The stiffness parameter.
+            D: The damping parameter of the soft contacts model.
+            mu: The static friction coefficient.
+            p:
+                The exponent p corresponding to the damping-related non-linearity
+                of the Hunt/Crossley model.
+            q:
+                The exponent q corresponding to the spring-related non-linearity
+                of the Hunt/Crossley model
+
+        Returns:
+            A tuple containing the computed contact force and the derivative of the
+            material deformation.
+        """
 
         # Convert the input vectors to arrays.
         W_p_C = jnp.array(position, dtype=float).squeeze()
@@ -364,6 +387,20 @@ class SoftContacts(common.ContactModel):
         parameters: SoftContactsParams,
         terrain: Terrain,
     ) -> tuple[jtp.Vector, jtp.Vector]:
+        """
+        Compute the contact force.
+
+        Args:
+            position: The position of the collidable point.
+            velocity: The velocity of the collidable point.
+            tangential_deformation: The material deformation of the collidable point.
+            parameters: The parameters of the soft contacts model.
+            terrain: The terrain model.
+
+        Returns:
+            A tuple containing the computed contact force and the derivative of the
+            material deformation.
+        """
 
         CW_fl, ṁ = SoftContacts.hunt_crossley_contact_model(
             position=position,

jaxsim/rbda/contacts/visco_elastic.py

@@ -206,6 +206,7 @@ class ViscoElasticContacts(common.ContactModel):
                 If passed, it is used to estimate good default parameters.
             max_squarings:
                 The maximum number of squarings performed in the matrix exponential.
+            **kwargs: Extra arguments to ignore.
 
         Returns:
             The `ViscoElasticContacts` instance.