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

jaxsim/math/rotation.py

@@ -8,6 +8,9 @@ from .utils import safe_norm
 
 
 class Rotation:
+    """
+    A utility class for rotation matrix operations.
+    """
 
     @staticmethod
     def x(theta: jtp.Float) -> jtp.Matrix:
@@ -15,7 +18,7 @@ class Rotation:
         Generate a 3D rotation matrix around the X-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.
@@ -29,7 +32,7 @@ class Rotation:
         Generate a 3D rotation matrix around the Y-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.
@@ -43,7 +46,7 @@ class Rotation:
         Generate a 3D rotation matrix around the Z-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
             jtp.Matrix: 3D rotation matrix.

jaxsim/math/skew.py

@@ -14,7 +14,7 @@ class Skew:
         Compute the skew-symmetric matrix (wedge operator) of a 3D vector.
 
         Args:
-            vector (jtp.Vector): A 3D vector.
+            vector: A 3D vector.
 
         Returns:
             jtp.Matrix: The skew-symmetric matrix corresponding to the input vector.
@@ -31,7 +31,7 @@ class Skew:
         Extract the 3D vector from a skew-symmetric matrix (vee operator).
 
         Args:
-            matrix (jtp.Matrix): A 3x3 skew-symmetric matrix.
+            matrix: A 3x3 skew-symmetric matrix.
 
         Returns:
             jtp.Vector: The 3D vector extracted from the input matrix.

jaxsim/math/transform.py

@@ -5,6 +5,9 @@ import jaxsim.typing as jtp
 
 
 class Transform:
+    """
+    A utility class for transformation matrix operations.
+    """
 
     @staticmethod
     def from_quaternion_and_translation(

jaxsim/math/utils.py

@@ -5,8 +5,8 @@ import jaxsim.typing as jtp
 
 def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
     """
-    Provides a calculation for an array norm so that it is safe
-    to compute the gradient and handle NaNs.
+    Compute an array norm handling NaNs and making sure that
+    it is safe to get the gradient.
 
     Args:
         array: The array for which to compute the norm.

jaxsim/mujoco/loaders.py

@@ -22,7 +22,7 @@ def load_rod_model(
     model_name: str | None = None,
 ) -> rod.Model:
     """
-    Loads a ROD model from a URDF/SDF file or a ROD model.
+    Load a ROD model from a URDF/SDF file or a ROD model.
 
     Args:
         model_description: The URDF/SDF file or ROD model to load.
@@ -62,14 +62,16 @@ def load_rod_model(
 
 
 class RodModelToMjcf:
-    """"""
+    """
+    Class to convert a ROD model to a Mujoco MJCF string.
+    """
 
     @staticmethod
     def assets_from_rod_model(
         rod_model: rod.Model,
     ) -> dict[str, bytes]:
         """
-        Generates a dictionary of assets from a ROD model.
+        Generate a dictionary of assets from a ROD model.
 
         Args:
             rod_model: The ROD model to extract the assets from.
@@ -112,7 +114,7 @@ class RodModelToMjcf:
         floating_joint_name: str = "world_to_base",
     ) -> str:
         """
-        Adds a floating joint to a URDF string.
+        Add a floating joint to a URDF string.
 
         Args:
             urdf_string: The URDF string to modify.
@@ -171,7 +173,7 @@ class RodModelToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a ROD model to a Mujoco MJCF string.
+        Convert a ROD model to a Mujoco MJCF string.
 
         Args:
             rod_model: The ROD model to convert.
@@ -522,6 +524,10 @@ class RodModelToMjcf:
 
 
 class UrdfToMjcf:
+    """
+    Class to convert a URDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         urdf: str | pathlib.Path,
@@ -532,7 +538,7 @@ class UrdfToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a URDF file to a Mujoco MJCF string.
+        Convert a URDF file to a Mujoco MJCF string.
 
         Args:
             urdf: The URDF file to convert.
@@ -564,6 +570,10 @@ class UrdfToMjcf:
 
 
 class SdfToMjcf:
+    """
+    Class to convert a SDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         sdf: str | pathlib.Path,
@@ -574,7 +584,7 @@ class SdfToMjcf:
         cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
-        Converts a SDF file to a Mujoco MJCF string.
+        Convert a SDF file to a Mujoco MJCF string.
 
         Args:
             sdf: The SDF file to convert.

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