jaxsim 0.2.1.dev69__py3-none-any.whl → 0.2.1.dev70__py3-none-any.whl

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/api/__init__.py

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

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/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/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.dist-info → jaxsim-0.2.1.dev70.dist-info}/METADATA

@@ -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.dist-info → jaxsim-0.2.1.dev70.dist-info}/RECORD

@@ -1,16 +1,17 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=xK6IWQvp5Q8AhpgxWdFhkwoBgV1BAj0kenm6gO6d_z4,426
+jaxsim/_version.py,sha256=1SWO53mcoalXzzBdA6_IrZRpZQsi_WIK6NLf67IAp-A,426
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
-jaxsim/api/__init__.py,sha256=fNTCPUeDfOAizRd4RsW3Epv0sLTu0KJGoFRSEsi75VM,162
+jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=Qtm_6qpiK4WtDVn6JMUHa8DySgBl9CjgKCybJqZ58Lc,7379
 jaxsim/api/common.py,sha256=DV-WZG28sikXopNv458aYvpLjmiAtFr5LRscOwXusuk,6640
 jaxsim/api/contact.py,sha256=Ve4ZOWkLEBRgK3KhtICxKY7YzsxYvc3lO-pPRBjqSnY,8659
 jaxsim/api/data.py,sha256=GXSCfq4_PsRZuUwsmcl18e2Ppdtu4ykYs0oWih8J2ZI,26775
+jaxsim/api/frame.py,sha256=68liCmwFDikuP0GQOpo-nF7aVeHJ173e1USSOwskTP0,6246
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
 jaxsim/api/kin_dyn_parameters.py,sha256=9xdHS3lDWSEdV0qAqCVoS6cSMUw17ebPJdbBfq4vjVU,25959
 jaxsim/api/link.py,sha256=rypTwkMf9HJ5UuAtHRJh0LqqdJWcLKTtTjWcjduEsF0,9842
-jaxsim/api/model.py,sha256=Ci7hl4M1gbnHNQAO5GLjxOdvJe8w3dhvGB_7x2gdpaU,53540
+jaxsim/api/model.py,sha256=7qBsneoc3dWNIsd6INzItGaIWSTgO98TQUZaeyVRrdk,53809
 jaxsim/api/ode.py,sha256=6l-6i2YHagsQvR8Ac-_fmO6P0hBVT6NkHhwXnrdITEg,9785
 jaxsim/api/ode_data.py,sha256=dwRFVPJ30XMmdUbPXEu7YxsQ97jZP4L4fd5ZzhrO5Ys,22184
 jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
@@ -37,7 +38,7 @@ jaxsim/parsers/kinematic_graph.py,sha256=zFt7x7pPGJar36Azukdi1eI_sa1kMWD3B8kZqcH
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
 jaxsim/parsers/descriptions/collision.py,sha256=HUWwuRgI9KznY29FFw1_zU3bGigDEezrcPOJSxSJGNU,3382
 jaxsim/parsers/descriptions/joint.py,sha256=r7teUWJyw62WS7qNCB_9FGQD1hvGHBepk-2LCSlmcSc,3052
-jaxsim/parsers/descriptions/link.py,sha256=hqLLitrAXnouia6ULk1BPOIEfRxrXwHmoPsi306IZW8,2859
+jaxsim/parsers/descriptions/link.py,sha256=GC-6ZgRZuRVpcRo1sY6YaR8lkCHkR4DvHNs2Ydw_tn4,2887
 jaxsim/parsers/descriptions/model.py,sha256=uO5xOJtViihVPnSSsmfQJvCh45ANyi9KYAzLOhH0R8g,8993
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
 jaxsim/parsers/rod/parser.py,sha256=Q13TOkmpU0SHpgSV8WRYWb290aPNNLsaz4eMlD4Mq5w,13525
@@ -57,8 +58,8 @@ jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
 jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim-0.2.1.dev69.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev69.dist-info/METADATA,sha256=UfQqQJzLiDAe6Mw8ZeHwUCawOIQ91K8IJ1MH6_sC_W4,9744
-jaxsim-0.2.1.dev69.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev69.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev69.dist-info/RECORD,,
+jaxsim-0.2.1.dev70.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev70.dist-info/METADATA,sha256=4AdUyv7PhRO09XpEBc_79fG0MC-msgEQ-CnrstJesqA,9744
+jaxsim-0.2.1.dev70.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev70.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev70.dist-info/RECORD,,