jaxsim 0.2.1.dev69__tar.gz → 0.2.1.dev70__tar.gz

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/.gitignore

@@ -1,5 +1,6 @@
 # IDEs
 .idea*
+.vscode/
 
 # Matlab
 *.m~
@@ -141,4 +142,3 @@ src/jaxsim/_version.py
 .ruff_cache/
 # pixi environments
 .pixi
-

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/PKG-INFO

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

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/docs/modules/api.rst

@@ -38,6 +38,11 @@ Link
 .. automodule:: jaxsim.api.link
     :members:
 
+Frame
+~~~~~
+.. automodule:: jaxsim.api.frame
+    :members:
+
 CoM
 ~~~
 .. automodule:: jaxsim.api.com

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/src/jaxsim/_version.py

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

jaxsim-0.2.1.dev70/src/jaxsim/api/__init__.py

@@ -0,0 +1,13 @@
+from . import common  # isort:skip
+from . import model, data  # isort:skip
+from . import (
+    com,
+    contact,
+    frame,
+    joint,
+    kin_dyn_parameters,
+    link,
+    ode,
+    ode_data,
+    references,
+)

jaxsim-0.2.1.dev70/src/jaxsim/api/frame.py

@@ -0,0 +1,221 @@
+import functools
+from typing import Sequence
+
+import jax
+import jax.numpy as jnp
+import jaxlie
+import numpy as np
+
+import jaxsim.api as js
+import jaxsim.math
+import jaxsim.typing as jtp
+
+from .common import VelRepr
+
+# =======================
+# Index-related functions
+# =======================
+
+
+def idx_of_parent_link(model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike) -> int:
+    """
+    Get the index of the link to which the frame is rigidly attached.
+
+    Args:
+        model: The model to consider.
+        frame_idx: The index of the frame.
+
+    Returns:
+        The index of the frame's parent link.
+    """
+
+    # Get the intermediate representation parsed from the model description.
+    ir = model.description.get()
+
+    # Extract the indices of the frame and the link it is attached to.
+    F = ir.frames[frame_idx - model.number_of_links()]
+    L = ir.links_dict[F.parent.name].index
+
+    return int(L)
+
+
+def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> int:
+    """
+    Convert the name of a frame to its index.
+
+    Args:
+        model: The model to consider.
+        frame_name: The name of the frame.
+
+    Returns:
+        The index of the frame.
+    """
+
+    frame_names = np.array([frame.name for frame in model.description.get().frames])
+
+    if frame_name in frame_names:
+        idx_in_list = np.argwhere(frame_names == frame_name)
+        return int(idx_in_list.squeeze().tolist()) + model.number_of_links()
+
+    return -1
+
+
+def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
+    """
+    Convert the index of a frame to its name.
+
+    Args:
+        model: The model to consider.
+        frame_index: The index of the frame.
+
+    Returns:
+        The name of the frame.
+    """
+
+    return model.description.get().frames[frame_index - model.number_of_links()].name
+
+
+@functools.partial(jax.jit, static_argnames=["frame_names"])
+def names_to_idxs(
+    model: js.model.JaxSimModel, *, frame_names: Sequence[str]
+) -> jax.Array:
+    """
+    Convert a sequence of frame names to their corresponding indices.
+
+    Args:
+        model: The model to consider.
+        frame_names: The names of the frames.
+
+    Returns:
+        The indices of the frames.
+    """
+
+    return jnp.array(
+        [name_to_idx(model=model, frame_name=frame_name) for frame_name in frame_names]
+    ).astype(int)
+
+
+def idxs_to_names(
+    model: js.model.JaxSimModel, *, frame_indices: Sequence[jtp.IntLike]
+) -> tuple[str, ...]:
+    """
+    Convert a sequence of frame indices to their corresponding names.
+
+    Args:
+        model: The model to consider.
+        frame_indices: The indices of the frames.
+
+    Returns:
+        The names of the frames.
+    """
+
+    return tuple(
+        idx_to_name(model=model, frame_index=frame_index)
+        for frame_index in frame_indices
+    )
+
+
+# ==========
+# Frame APIs
+# ==========
+
+
+@functools.partial(jax.jit, static_argnames=["frame_index"])
+def transform(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+) -> jtp.Matrix:
+    """
+    Compute the SE(3) transform from the world frame to the specified frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame for which the transform is requested.
+
+    Returns:
+        The 4x4 matrix representing the transform.
+    """
+
+    # Compute the necessary transforms.
+    L = idx_of_parent_link(model=model, frame_idx=frame_index)
+    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()]
+    L_H_F = frame.pose
+
+    # Combine the transforms computing the frame pose.
+    return W_H_L @ L_H_F
+
+
+@functools.partial(jax.jit, static_argnames=["frame_index", "output_vel_repr"])
+def jacobian(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the free-floating jacobian of the frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian.
+
+    Returns:
+        The 6×(6+n) free-floating jacobian of the frame.
+
+    Note:
+        The input representation of the free-floating jacobian is the active
+        velocity representation.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the index of the parent link.
+    L = idx_of_parent_link(model=model, frame_idx=frame_index)
+
+    # Compute the Jacobian of the parent link using body-fixed output representation.
+    L_J_WL = js.link.jacobian(
+        model=model, data=data, link_index=L, output_vel_repr=VelRepr.Body
+    )
+
+    # Adjust the output representation
+    match output_vel_repr:
+        case VelRepr.Inertial:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_X_L = jaxlie.SE3.from_matrix(W_H_L).adjoint()
+            W_J_WL = W_X_L @ L_J_WL
+            O_J_WL_I = W_J_WL
+
+        case VelRepr.Body:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            F_H_L = jaxsim.math.Transform.inverse(W_H_F) @ W_H_L
+            F_X_L = jaxlie.SE3.from_matrix(F_H_L).adjoint()
+            F_J_WL = F_X_L @ L_J_WL
+            O_J_WL_I = F_J_WL
+
+        case VelRepr.Mixed:
+            W_H_L = js.link.transform(model=model, data=data, link_index=L)
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            F_H_L = jaxsim.math.Transform.inverse(W_H_F) @ W_H_L
+            FW_H_F = W_H_F.at[0:3, 3].set(jnp.zeros(3))
+            FW_H_L = FW_H_F @ F_H_L
+            FW_X_L = jaxlie.SE3.from_matrix(FW_H_L).adjoint()
+            FW_J_WL = FW_X_L @ L_J_WL
+            O_J_WL_I = FW_J_WL
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    return O_J_WL_I

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/src/jaxsim/api/model.py

@@ -252,6 +252,16 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.link_names
 
+    def frame_names(self) -> tuple[str, ...]:
+        """
+        Return the names of the links in the model.
+
+        Returns:
+            The names of the links in the model.
+        """
+
+        return tuple([frame.name for frame in self.description.get().frames])
+
 
 # =====================
 # Model post-processing

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/src/jaxsim/parsers/descriptions/link.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 import dataclasses
-from typing import List
 
 import jax.numpy as jnp
 import jax_dataclasses
@@ -16,39 +17,46 @@ class LinkDescription(JaxsimDataclass):
     In-memory description of a robot link.
 
     Attributes:
-        name (str): The name of the link.
-        mass (float): The mass of the link.
-        inertia (jtp.Matrix): The inertia matrix of the link.
-        index (Optional[int]): An optional index for the link.
-        parent (Optional[LinkDescription]): The parent link of this link.
-        pose (jtp.Matrix): The pose transformation matrix of the link.
-        children (List[LinkDescription]): List of child links.
+        name: The name of the link.
+        mass: The mass of the link.
+        inertia: The inertia tensor of the link.
+        index: An optional index for the link (it gets automatically assigned).
+        parent: The parent link of this link.
+        pose: The pose transformation matrix of the link.
+        children: List of child links.
     """
 
     name: Static[str]
-    mass: float
-    inertia: jtp.Matrix
+    mass: float = dataclasses.field(repr=False)
+    inertia: jtp.Matrix = dataclasses.field(repr=False)
     index: int | None = None
-    parent: Static["LinkDescription"] = dataclasses.field(default=None, repr=False)
+    parent: LinkDescription = dataclasses.field(default=None, repr=False)
     pose: jtp.Matrix = dataclasses.field(default_factory=lambda: jnp.eye(4), repr=False)
-    children: Static[List["LinkDescription"]] = dataclasses.field(
+
+    children: Static[list[LinkDescription]] = dataclasses.field(
         default_factory=list, repr=False
     )
 
     def __hash__(self) -> int:
-        return hash(self.__repr__())
-
-    def __eq__(self, other) -> bool:
-        return (
-            self.name == other.name
-            and self.mass == other.mass
-            and (self.inertia == other.inertia).all()
-            and self.index == other.index
-            and self.parent == other.parent
-            and (self.pose == other.pose).all()
-            and self.children == other.children
+
+        return hash(
+            (
+                hash(self.name),
+                hash(float(self.mass)),
+                hash(tuple(self.inertia.flatten().tolist())),
+                hash(int(self.index)),
+                hash(self.parent),
+                hash(tuple(hash(c) for c in self.children)),
+            )
         )
 
+    def __eq__(self, other: LinkDescription) -> bool:
+
+        if not isinstance(other, LinkDescription):
+            return False
+
+        return hash(self) == hash(other)
+
     @property
     def name_and_index(self) -> str:
         """
@@ -61,19 +69,19 @@ class LinkDescription(JaxsimDataclass):
         return f"#{self.index}_<{self.name}>"
 
     def lump_with(
-        self, link: "LinkDescription", lumped_H_removed: jtp.Matrix
-    ) -> "LinkDescription":
+        self, link: LinkDescription, lumped_H_removed: jtp.Matrix
+    ) -> LinkDescription:
         """
         Combine the current link with another link, preserving mass and inertia.
 
         Args:
-            link (LinkDescription): The link to combine with.
-            lumped_H_removed (jtp.Matrix): The transformation matrix between the two links.
+            link: The link to combine with.
+            lumped_H_removed: The transformation matrix between the two links.
 
         Returns:
-            LinkDescription: The combined link.
-
+            The combined link.
         """
+
         # Get the 6D inertia of the link to remove
         I_removed = link.inertia
 

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/src/jaxsim.egg-info/PKG-INFO

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

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/src/jaxsim.egg-info/SOURCES.txt

@@ -51,6 +51,7 @@ src/jaxsim/api/com.py
 src/jaxsim/api/common.py
 src/jaxsim/api/contact.py
 src/jaxsim/api/data.py
+src/jaxsim/api/frame.py
 src/jaxsim/api/joint.py
 src/jaxsim/api/kin_dyn_parameters.py
 src/jaxsim/api/link.py
@@ -105,6 +106,7 @@ tests/__init__.py
 tests/conftest.py
 tests/test_api_com.py
 tests/test_api_data.py
+tests/test_api_frame.py
 tests/test_api_joint.py
 tests/test_api_link.py
 tests/test_api_model.py

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/tests/conftest.py

@@ -4,6 +4,8 @@ import pathlib
 import jax
 import pytest
 import rod
+import rod.urdf
+import rod.urdf.exporter
 
 import jaxsim
 import jaxsim.api as js
@@ -120,18 +122,26 @@ def jaxsim_model_box() -> js.model.JaxSimModel:
     rod_model = (
         rod.builder.primitives.BoxBuilder(x=0.3, y=0.2, z=0.1, mass=1.0, name="box")
         .build_model()
-        .add_link()
+        .add_link(name="box_link")
         .add_inertial()
         .add_visual()
         .add_collision()
         .build()
     )
 
-    # Export the URDF string.
-    urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
-        sdf=rod_model, pretty=True
+    rod_model.add_frame(
+        rod.Frame(
+            name="box_frame",
+            attached_to="box_link",
+            pose=rod.Pose(relative_to="box_link", pose=[1, 1, 1, 0.5, 0.4, 0.3]),
+        )
     )
 
+    # Export the URDF string.
+    urdf_string = rod.urdf.exporter.UrdfExporter(
+        pretty=True, gazebo_preserve_fixed_joints=True
+    ).to_urdf_string(sdf=rod_model)
+
     return build_jaxsim_model(model_description=urdf_string)
 
 
@@ -159,8 +169,8 @@ def jaxsim_model_sphere() -> js.model.JaxSimModel:
     )
 
     # Export the URDF string.
-    urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
-        sdf=rod_model, pretty=True
+    urdf_string = rod.urdf.exporter.UrdfExporter(pretty=True).to_urdf_string(
+        sdf=rod_model
     )
 
     return build_jaxsim_model(model_description=urdf_string)
@@ -219,7 +229,9 @@ def jaxsim_model_ergocub_reduced(jaxsim_model_ergocub) -> js.model.JaxSimModel:
         and "torso" not in j and "elbow" not in j and "shoulder" not in j
     )
 
-    return js.model.reduce(model=model_full, considered_joints=reduced_joints)
+    model = js.model.reduce(model=model_full, considered_joints=reduced_joints)
+
+    return model
 
 
 @pytest.fixture(scope="session")

jaxsim-0.2.1.dev70/tests/test_api_frame.py

@@ -0,0 +1,144 @@
+import jax
+import numpy as np
+import pytest
+
+import jaxsim.api as js
+from jaxsim import VelRepr
+
+from . import utils_idyntree
+
+
+def test_frame_index(jaxsim_models_types: js.model.JaxSimModel):
+
+    model = jaxsim_models_types
+
+    # =====
+    # Tests
+    # =====
+
+    frame_indices = tuple(
+        frame.index
+        for frame in model.description.get().frames
+        if frame.index is not None
+    )
+
+    frame_names = np.array([frame.name for frame in model.description.get().frames])
+
+    for frame_idx, frame_name in zip(frame_indices, frame_names):
+        assert js.frame.name_to_idx(model=model, frame_name=frame_name) == frame_idx
+        assert js.frame.idx_to_name(model=model, frame_index=frame_idx) == frame_name
+        assert (
+            js.frame.idx_of_parent_link(model=model, frame_idx=frame_idx)
+            < model.number_of_links()
+        )
+
+    assert js.frame.names_to_idxs(
+        model=model, frame_names=tuple(frame_names)
+    ) == pytest.approx(frame_indices)
+
+    assert js.frame.idxs_to_names(
+        model=model, frame_indices=frame_indices
+    ) == pytest.approx(frame_names)
+
+
+def test_frame_transforms(
+    jaxsim_models_types: js.model.JaxSimModel,
+    prng_key: jax.Array,
+):
+
+    model = jaxsim_models_types
+
+    key, subkey = jax.random.split(prng_key, num=2)
+    data = js.data.random_model_data(
+        model=model, key=subkey, velocity_representation=VelRepr.Inertial
+    )
+
+    kin_dyn = utils_idyntree.build_kindyncomputations_from_jaxsim_model(
+        model=model, data=data
+    )
+
+    # Get all names of frames in the iDynTree model.
+    frame_names = [
+        frame.name
+        for frame in model.description.get().frames
+        if frame.name in kin_dyn.frame_names()
+    ]
+
+    # Skip some entry of models with many frames.
+    frame_names = [
+        name
+        for name in frame_names
+        if "skin" not in name or "laser" not in name or "depth" not in name
+    ]
+
+    # Get indices of frames.
+    frame_indices = tuple(
+        frame.index
+        for frame in model.description.get().frames
+        if frame.index is not None and frame.name in frame_names
+    )
+
+    # =====
+    # Tests
+    # =====
+
+    assert len(frame_indices) == len(frame_names)
+
+    for frame_name, frame_idx in zip(frame_names, frame_indices):
+
+        W_H_F_js = js.frame.transform(
+            model=model,
+            data=data,
+            frame_index=js.frame.name_to_idx(model=model, frame_name=frame_name),
+        )
+        W_H_F_idt = kin_dyn.frame_transform(frame_name=frame_name)
+        assert W_H_F_js == pytest.approx(W_H_F_idt, abs=1e-6), frame_name
+
+
+def test_frame_jacobians(
+    jaxsim_models_types: js.model.JaxSimModel,
+    velocity_representation: VelRepr,
+    prng_key: jax.Array,
+):
+
+    model = jaxsim_models_types
+
+    key, subkey = jax.random.split(prng_key, num=2)
+    data = js.data.random_model_data(
+        model=model, key=subkey, velocity_representation=velocity_representation
+    )
+
+    kin_dyn = utils_idyntree.build_kindyncomputations_from_jaxsim_model(
+        model=model, data=data
+    )
+
+    # Get all names of frames in the iDynTree model.
+    frame_names = [
+        frame.name
+        for frame in model.description.get().frames
+        if frame.name in kin_dyn.frame_names()
+    ]
+
+    # Lower the number of frames for models with many frames.
+    if model.name().lower() == "ergocub":
+        assert any(["sole" in name for name in frame_names])
+        frame_names = [name for name in frame_names if "sole" in name]
+
+    # Get indices of frames.
+    frame_indices = tuple(
+        frame.index
+        for frame in model.description.get().frames
+        if frame.index is not None and frame.name in frame_names
+    )
+
+    # =====
+    # Tests
+    # =====
+
+    assert len(frame_indices) == len(frame_names)
+
+    for frame_name, frame_idx in zip(frame_names, frame_indices):
+
+        J_WL_js = js.frame.jacobian(model=model, data=data, frame_index=frame_idx)
+        J_WL_idt = kin_dyn.jacobian_frame(frame_name=frame_name)
+        assert J_WL_js == pytest.approx(J_WL_idt, abs=1e-9)

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/tests/test_api_joint.py

@@ -16,6 +16,7 @@ def test_joint_index(
 
     for idx, joint_name in enumerate(model.joint_names()):
         assert js.joint.name_to_idx(model=model, joint_name=joint_name) == idx
+        assert js.joint.idx_to_name(model=model, joint_index=idx) == joint_name
 
     assert js.joint.names_to_idxs(
         model=model, joint_names=model.joint_names()

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/tests/test_api_link.py

@@ -20,6 +20,7 @@ def test_link_index(
 
     for idx, link_name in enumerate(model.link_names()):
         assert js.link.name_to_idx(model=model, link_name=link_name) == idx
+        assert js.link.idx_to_name(model=model, link_index=idx) == link_name
 
     assert js.link.names_to_idxs(
         model=model, link_names=model.link_names()

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/tests/test_pytree.py

@@ -2,6 +2,7 @@ import io
 from contextlib import redirect_stdout
 
 import jax
+import pytest
 import rod.builder.primitives
 import rod.urdf.exporter
 
@@ -9,6 +10,7 @@ import jaxsim.api as js
 
 
 # https://github.com/ami-iit/jaxsim/issues/103
+@pytest.mark.xfail(strict=True)
 def test_call_jit_compiled_function_passing_different_objects():
 
     # Create on-the-fly a ROD model of a box.
@@ -23,8 +25,8 @@ def test_call_jit_compiled_function_passing_different_objects():
     )
 
     # Export the URDF string.
-    urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
-        sdf=rod_model, pretty=True
+    urdf_string = rod.urdf.exporter.UrdfExporter(pretty=True).to_urdf_string(
+        sdf=rod_model
     )
 
     model1 = js.model.JaxSimModel.build_from_model_description(

{jaxsim-0.2.1.dev69 → jaxsim-0.2.1.dev70}/tests/utils_idyntree.py

@@ -348,6 +348,13 @@ class KinDynComputations:
 
         return H
 
+    def frame_parent_link_name(self, frame_name: str) -> str:
+        return self.kin_dyn.model().getLinkName(
+            self.kin_dyn.model().getFrameLink(
+                self.kin_dyn.model().getFrameIndex(frame_name)
+            )
+        )
+
     def base_velocity(self) -> npt.NDArray:
 
         nu = idt.VectorDynSize()

jaxsim-0.2.1.dev69/src/jaxsim/api/__init__.py

@@ -1,3 +0,0 @@
-from . import common  # isort:skip
-from . import model, data  # isort:skip
-from . import com, contact, joint, kin_dyn_parameters, link, ode, ode_data, references