jaxsim 0.5.1.dev48__py3-none-any.whl → 0.5.1.dev60__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev48'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev48')
+__version__ = version = '0.5.1.dev60'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev60')

jaxsim/api/com.py

@@ -8,6 +8,7 @@ import jaxsim.typing as jtp
 from .common import VelRepr
 
 
+@js.common.named_scope
 @jax.jit
 def com_position(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -44,6 +45,7 @@ def com_position(
     return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def com_linear_velocity(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -74,6 +76,7 @@ def com_linear_velocity(
     return G_vl_WG
 
 
+@js.common.named_scope
 @jax.jit
 def centroidal_momentum(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -101,6 +104,7 @@ def centroidal_momentum(
     return G_J @ ν
 
 
+@js.common.named_scope
 @jax.jit
 def centroidal_momentum_jacobian(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -149,6 +153,7 @@ def centroidal_momentum_jacobian(
     return G_Xf_B @ B_Jh
 
 
+@js.common.named_scope
 @jax.jit
 def locked_centroidal_spatial_inertia(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -186,6 +191,7 @@ def locked_centroidal_spatial_inertia(
     return G_Xf_B @ B_Mbb_B @ B_Xv_G
 
 
+@js.common.named_scope
 @jax.jit
 def average_centroidal_velocity(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -213,6 +219,7 @@ def average_centroidal_velocity(
     return G_J @ ν
 
 
+@js.common.named_scope
 @jax.jit
 def average_centroidal_velocity_jacobian(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -239,6 +246,7 @@ def average_centroidal_velocity_jacobian(
     return jnp.linalg.inv(G_Mbb) @ G_J
 
 
+@js.common.named_scope
 @jax.jit
 def bias_acceleration(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData

jaxsim/api/common.py

@@ -3,7 +3,8 @@ import contextlib
 import dataclasses
 import enum
 import functools
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
+from typing import ParamSpec, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -20,6 +21,21 @@ except ImportError:
     from typing_extensions import Self
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
+    """Applies a JAX named scope to a function for improved profiling and clarity."""
+
+    @functools.wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        with jax.named_scope(name or fn.__name__):
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
 @enum.unique
 class VelRepr(enum.IntEnum):
     """

jaxsim/api/contact.py

@@ -16,6 +16,7 @@ from jaxsim.rbda import contacts
 from .common import VelRepr
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_kinematics(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -52,6 +53,7 @@ def collidable_point_kinematics(
     return W_p_Ci, W_ṗ_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_positions(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -72,6 +74,7 @@ def collidable_point_positions(
     return W_p_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_velocities(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -92,6 +95,7 @@ def collidable_point_velocities(
     return W_ṗ_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_forces(
     model: js.model.JaxSimModel,
@@ -129,6 +133,7 @@ def collidable_point_forces(
     return f_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_dynamics(
     model: js.model.JaxSimModel,
@@ -227,6 +232,7 @@ def collidable_point_dynamics(
     return f_Ci, aux_data
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["link_names"])
 def in_contact(
     model: js.model.JaxSimModel,
@@ -424,6 +430,7 @@ def estimate_good_contact_parameters(
     return parameters
 
 
+@js.common.named_scope
 @jax.jit
 def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     r"""
@@ -469,6 +476,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
@@ -561,6 +569,7 @@ def jacobian(
     return O_J_WC
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian_derivative(
     model: js.model.JaxSimModel,

jaxsim/api/data.py

@@ -253,6 +253,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return -self.gravity[2]
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_positions(
         self,
@@ -300,6 +301,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.joint_positions[joint_idxs]
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_velocities(
         self,
@@ -347,6 +349,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.joint_velocities[joint_idxs]
 
+    @js.common.named_scope
     @jax.jit
     def base_position(self) -> jtp.Vector:
         """
@@ -358,6 +361,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.base_position.squeeze()
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["dcm"])
     def base_orientation(self, dcm: jtp.BoolLike = False) -> jtp.Vector | jtp.Matrix:
         """
@@ -386,6 +390,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             float
         )
 
+    @js.common.named_scope
     @jax.jit
     def base_transform(self) -> jtp.Matrix:
         """
@@ -405,6 +410,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ]
         )
 
+    @js.common.named_scope
     @jax.jit
     def base_velocity(self) -> jtp.Vector:
         """
@@ -434,6 +440,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             .astype(float)
         )
 
+    @js.common.named_scope
     @jax.jit
     def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
         r"""
@@ -446,6 +453,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.base_transform(), self.joint_positions()
 
+    @js.common.named_scope
     @jax.jit
     def generalized_velocity(self) -> jtp.Vector:
         r"""
@@ -466,6 +474,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Store quantities
     # ================
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def reset_joint_positions(
         self,
@@ -514,6 +523,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             s=self.state.physics_model.joint_positions.at[joint_idxs].set(positions)
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def reset_joint_velocities(
         self,
@@ -562,6 +572,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ṡ=self.state.physics_model.joint_velocities.at[joint_idxs].set(velocities)
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_position(self, base_position: jtp.VectorLike) -> Self:
         """
@@ -585,6 +596,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
         """
@@ -612,6 +624,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_pose(self, base_pose: jtp.MatrixLike) -> Self:
         """
@@ -634,6 +647,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_quaternion=W_Q_B
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_linear_velocity(
         self,
@@ -665,6 +679,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_angular_velocity(
         self,
@@ -696,6 +711,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_velocity(
         self,

jaxsim/api/frame.py

@@ -16,6 +16,7 @@ from .common import VelRepr
 # =======================
 
 
+@js.common.named_scope
 @jax.jit
 def idx_of_parent_link(
     model: js.model.JaxSimModel, *, frame_index: jtp.IntLike
@@ -45,6 +46,7 @@ def idx_of_parent_link(
     ]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="frame_name")
 def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
     """
@@ -97,6 +99,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str
     ]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["frame_names"])
 def names_to_idxs(
     model: js.model.JaxSimModel, *, frame_names: Sequence[str]
@@ -139,6 +142,7 @@ def idxs_to_names(
 # ==========
 
 
+@js.common.named_scope
 @jax.jit
 def transform(
     model: js.model.JaxSimModel,
@@ -180,6 +184,7 @@ def transform(
     return W_H_L @ L_H_F
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def velocity(
     model: js.model.JaxSimModel,
@@ -230,6 +235,7 @@ def velocity(
     return O_J_WF_I @ I_ν
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
@@ -309,6 +315,7 @@ def jacobian(
     return O_J_WL_I
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian_derivative(
     model: js.model.JaxSimModel,

jaxsim/api/joint.py

@@ -13,6 +13,7 @@ from jaxsim import exceptions
 # =======================
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="joint_name")
 def name_to_idx(model: js.model.JaxSimModel, *, joint_name: str) -> jtp.Int:
     """
@@ -61,6 +62,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
     return model.kin_dyn_parameters.joint_model.joint_names[joint_index + 1]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="joint_names")
 def names_to_idxs(
     model: js.model.JaxSimModel, *, joint_names: Sequence[str]
@@ -141,6 +143,7 @@ def position_limit(
     return s_min.astype(float), s_max.astype(float)
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["joint_names"])
 def position_limits(
     model: js.model.JaxSimModel, *, joint_names: Sequence[str] | None = None
@@ -176,6 +179,7 @@ def position_limits(
 # ======================
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["joint_names"])
 def random_joint_positions(
     model: js.model.JaxSimModel,

jaxsim/api/model.py

@@ -491,6 +491,7 @@ def reduce(
 # ===================
 
 
+@js.common.named_scope
 @jax.jit
 def total_mass(model: JaxSimModel) -> jtp.Float:
     """
@@ -506,6 +507,7 @@ def total_mass(model: JaxSimModel) -> jtp.Float:
     return model.kin_dyn_parameters.link_parameters.mass.sum().astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def link_spatial_inertia_matrices(model: JaxSimModel) -> jtp.Array:
     """
@@ -528,6 +530,7 @@ def link_spatial_inertia_matrices(model: JaxSimModel) -> jtp.Array:
 # ==============================
 
 
+@js.common.named_scope
 @jax.jit
 def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     """
@@ -915,6 +918,7 @@ def forward_dynamics(
     )
 
 
+@js.common.named_scope
 @jax.jit
 def forward_dynamics_aba(
     model: JaxSimModel,
@@ -1059,6 +1063,7 @@ def forward_dynamics_aba(
     return C_v̇_WB.astype(float), s̈.astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def forward_dynamics_crb(
     model: JaxSimModel,
@@ -1150,6 +1155,7 @@ def forward_dynamics_crb(
     return v̇_WB, s̈
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_mass_matrix(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1195,6 +1201,7 @@ def free_floating_mass_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_coriolis_matrix(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1311,6 +1318,7 @@ def free_floating_coriolis_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@js.common.named_scope
 @jax.jit
 def inverse_dynamics(
     model: JaxSimModel,
@@ -1466,6 +1474,7 @@ def inverse_dynamics(
     return f_B.astype(float), τ.astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_gravity_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1515,6 +1524,7 @@ def free_floating_gravity_forces(
     ).astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_bias_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1584,6 +1594,7 @@ def free_floating_bias_forces(
 # ==========================
 
 
+@js.common.named_scope
 @jax.jit
 def locked_spatial_inertia(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1602,6 +1613,7 @@ def locked_spatial_inertia(
     return total_momentum_jacobian(model=model, data=data)[:, 0:6]
 
 
+@js.common.named_scope
 @jax.jit
 def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
     """
@@ -1690,6 +1702,7 @@ def total_momentum_jacobian(
             raise ValueError(output_vel_repr)
 
 
+@js.common.named_scope
 @jax.jit
 def average_velocity(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
     """
@@ -1778,6 +1791,7 @@ def average_velocity_jacobian(
 # ========================
 
 
+@js.common.named_scope
 @jax.jit
 def link_bias_accelerations(
     model: JaxSimModel,
@@ -1987,6 +2001,7 @@ def link_bias_accelerations(
     return O_v̇_WL
 
 
+@js.common.named_scope
 @jax.jit
 def link_contact_forces(
     model: js.model.JaxSimModel,
@@ -2062,6 +2077,7 @@ def link_contact_forces(
 # ======
 
 
+@js.common.named_scope
 @jax.jit
 def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2081,6 +2097,7 @@ def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.
     return (K + U).astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2102,6 +2119,7 @@ def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Flo
     return K.squeeze().astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2128,6 +2146,7 @@ def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.F
 # ==========
 
 
+@js.common.named_scope
 @jax.jit
 def step(
     model: JaxSimModel,
@@ -2198,7 +2217,9 @@ def step(
             isinstance(model.contact_model, jaxsim.rbda.contacts.ViscoElasticContacts)
             & (
                 ~jnp.allclose(dt, model.time_step)
-                | ~isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+                | ~int(
+                    isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+                )
             )
         ),
         msg=msg.format(module, name),

jaxsim/api/ode.py

@@ -85,6 +85,7 @@ def wrap_system_dynamics_for_integration(
 # ==================================
 
 
+@js.common.named_scope
 @jax.jit
 def system_velocity_dynamics(
     model: js.model.JaxSimModel,
@@ -331,6 +332,7 @@ def system_acceleration(
     return v̇_WB, s̈
 
 
+@js.common.named_scope
 @jax.jit
 def system_position_dynamics(
     model: js.model.JaxSimModel,
@@ -370,6 +372,7 @@ def system_position_dynamics(
     return W_ṗ_B, W_Q̇_B, ṡ
 
 
+@js.common.named_scope
 @jax.jit
 def system_dynamics(
     model: js.model.JaxSimModel,

jaxsim/api/ode_data.py

@@ -10,108 +10,14 @@ import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
 
-# =============================================================================
-# Define the input and state of the ODE system defining the integrated dynamics
-# =============================================================================
+# ===================================================================
+# Define the state of the ODE system defining the integrated dynamics
+# ===================================================================
 
 # Note: the ODE system is the combination of the floating-base dynamics and the
 #       soft-contacts dynamics.
 
 
-@jax_dataclasses.pytree_dataclass
-class ODEInput(JaxsimDataclass):
-    """
-    The input to the ODE system.
-
-    Attributes:
-        physics_model: The input to the physics model.
-    """
-
-    physics_model: PhysicsModelInput
-
-    @staticmethod
-    def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-    ) -> ODEInput:
-        """
-        Build an `ODEInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the ODE input.
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-
-        Returns:
-            The `ODEInput` built from the `JaxSimModel`.
-
-        Note:
-            If any of the input components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return ODEInput.build(
-            physics_model_input=PhysicsModelInput.build_from_jaxsim_model(
-                model=model,
-                link_forces=link_forces,
-                joint_force_references=joint_force_references,
-            ),
-            model=model,
-        )
-
-    @staticmethod
-    def build(
-        physics_model_input: PhysicsModelInput | None = None,
-        model: js.model.JaxSimModel | None = None,
-    ) -> ODEInput:
-        """
-        Build an `ODEInput` from a `PhysicsModelInput`.
-
-        Args:
-            physics_model_input: The `PhysicsModelInput` associated with the ODE input.
-            model: The `JaxSimModel` associated with the ODE input.
-
-        Returns:
-            A `ODEInput` instance.
-        """
-
-        physics_model_input = (
-            physics_model_input
-            if physics_model_input is not None
-            else PhysicsModelInput.zero(model=model)
-        )
-
-        return ODEInput(physics_model=physics_model_input)
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> ODEInput:
-        """
-        Build a zero `ODEInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the ODE input.
-
-        Returns:
-            A zero `ODEInput` instance.
-        """
-
-        return ODEInput.build(model=model)
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        """
-        Check if the `ODEInput` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `ODEInput` against.
-
-        Returns:
-            `True` if the ODE input is valid for the given model, `False` otherwise.
-        """
-
-        return self.physics_model.valid(model=model)
-
-
 @jax_dataclasses.pytree_dataclass
 class ODEState(JaxsimDataclass):
     """
@@ -493,123 +399,3 @@ class PhysicsModelState(JaxsimDataclass):
             return False
 
         return True
-
-
-@jax_dataclasses.pytree_dataclass
-class PhysicsModelInput(JaxsimDataclass):
-    """
-    Class storing the inputs of the physics model dynamics.
-
-    Attributes:
-        tau: The vector of joint forces.
-        f_ext: The matrix of external forces applied to the links.
-    """
-
-    tau: jtp.Vector
-    f_ext: jtp.Matrix
-
-    @staticmethod
-    def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-    ) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the input.
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-
-        Note:
-            If any of the input components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return PhysicsModelInput.build(
-            joint_force_references=joint_force_references,
-            link_forces=link_forces,
-            number_of_dofs=model.dofs(),
-            number_of_links=model.number_of_links(),
-        )
-
-    @staticmethod
-    def build(
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-        number_of_dofs: jtp.Int | None = None,
-        number_of_links: jtp.Int | None = None,
-    ) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput`.
-
-        Args:
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-            number_of_dofs: The number of degrees of freedom of the model.
-            number_of_links: The number of links of the model.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-        """
-
-        joint_force_references = jnp.atleast_1d(
-            jnp.array(joint_force_references, dtype=float).squeeze()
-            if joint_force_references is not None
-            else jnp.zeros(number_of_dofs)
-        ).astype(float)
-
-        link_forces = jnp.atleast_2d(
-            jnp.array(link_forces, dtype=float).squeeze()
-            if link_forces is not None
-            else jnp.zeros(shape=(number_of_links, 6))
-        ).astype(float)
-
-        return PhysicsModelInput(
-            tau=joint_force_references,
-            f_ext=link_forces,
-        )
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput` with all components initialized to zero.
-
-        Args:
-            model: The `JaxSimModel` associated with the input.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-        """
-
-        return PhysicsModelInput.build_from_jaxsim_model(model=model)
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        """
-        Check if the `PhysicsModelInput` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `PhysicsModelInput` against.
-
-        Returns:
-            `True` if the `PhysicsModelInput` is valid for the given model,
-            `False` otherwise.
-        """
-
-        shape = self.tau.shape
-        expected_shape = (model.dofs(),)
-
-        if shape != expected_shape:
-            return False
-
-        shape = self.f_ext.shape
-        expected_shape = (model.number_of_links(), 6)
-
-        if shape != expected_shape:
-            return False
-
-        return True

jaxsim/api/references.py

@@ -12,7 +12,6 @@ from jaxsim import exceptions
 from jaxsim.utils.tracing import not_tracing
 
 from .common import VelRepr
-from .ode_data import ODEInput
 
 try:
     from typing import Self
@@ -24,9 +23,14 @@ except ImportError:
 class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     """
     Class containing the references for a `JaxSimModel` object.
+
+    Attributes:
+        _link_forces: The link 6D forces in inertial-fixed representation.
+        _joint_force_references: The joint force references.
     """
 
-    input: ODEInput
+    _link_forces: jtp.Matrix
+    _joint_force_references: jtp.Vector
 
     @staticmethod
     def zero(
@@ -94,17 +98,21 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         velocity_representation = (
             velocity_representation
             if velocity_representation is not None
-            else (
-                data.velocity_representation if data is not None else VelRepr.Inertial
-            )
+            else getattr(data, "velocity_representation", VelRepr.Inertial)
         )
 
         # Create a zero references object.
         references = JaxSimModelReferences(
-            input=ODEInput.zero(model=model),
+            _link_forces=f_L,
+            _joint_force_references=joint_force_references,
             velocity_representation=velocity_representation,
         )
 
+        # If the velocity representation is inertial-fixed, we can return
+        # the references directly, as we store the link forces in this frame.
+        if velocity_representation is VelRepr.Inertial:
+            return references
+
         # Store the joint force references.
         references = references.set_joint_force_references(
             forces=joint_force_references,
@@ -135,17 +143,27 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             `False` otherwise.
         """
 
-        valid = True
+        if model is None:
+            return True
+
+        shape = self._joint_force_references.shape
+        expected_shape = (model.dofs(),)
+
+        if shape != expected_shape:
+            return False
+
+        shape = self._link_forces.shape
+        expected_shape = (model.number_of_links(), 6)
 
-        if model is not None:
-            valid = valid and self.input.valid(model=model)
+        if shape != expected_shape:
+            return False
 
-        return valid
+        return True
 
     # ==================
     # Extract quantities
     # ==================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names"])
     def link_forces(
         self,
@@ -178,7 +196,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             e.g. to the contact model and other kinematic constraints.
         """
 
-        W_f_L = self.input.physics_model.f_ext
+        W_f_L = self._link_forces
 
         # Return all link forces in inertial-fixed representation using the implicit
         # serialization.
@@ -190,7 +208,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             if link_names is not None:
                 raise ValueError("Link names cannot be provided without a model")
 
-            return self.input.physics_model.f_ext
+            return W_f_L
 
         # If we have the model, we can extract the link names, if not provided.
         link_idxs = (
@@ -207,7 +225,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not_tracing(self.input.physics_model.f_ext) and not data.valid(model=model):
+        if not_tracing(self._link_forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
         # Helper function to convert a single 6D force to the active representation
@@ -264,9 +282,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             if joint_names is not None:
                 raise ValueError("Joint names cannot be provided without a model")
 
-            return self.input.physics_model.tau
+            return self._joint_force_references
 
-        if not_tracing(self.input.physics_model.tau) and not self.valid(model=model):
+        if not_tracing(self._joint_force_references) and not self.valid(model=model):
             msg = "The actuation object is not compatible with the provided model"
             raise ValueError(msg)
 
@@ -277,13 +295,13 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         )
 
         return jnp.atleast_1d(
-            self.input.physics_model.tau[joint_idxs].squeeze()
+            self._joint_force_references[joint_idxs].squeeze()
         ).astype(float)
 
     # ================
     # Store quantities
     # ================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def set_joint_force_references(
         self,
@@ -310,11 +328,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         def replace(forces: jtp.Vector) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
-                input=self.input.replace(
-                    physics_model=self.input.physics_model.replace(
-                        tau=jnp.atleast_1d(forces.squeeze()).astype(float)
-                    )
-                ),
+                _joint_force_references=jnp.atleast_1d(forces.squeeze()).astype(float),
             )
 
         if model is None:
@@ -330,8 +344,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             else jnp.arange(model.number_of_joints())
         )
 
-        return replace(forces=self.input.physics_model.tau.at[joint_idxs].set(forces))
+        return replace(forces=self._joint_force_references.at[joint_idxs].set(forces))
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names", "additive"])
     def apply_link_forces(
         self,
@@ -370,11 +385,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         def replace(forces: jtp.MatrixLike) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
-                input=self.input.replace(
-                    physics_model=self.input.physics_model.replace(
-                        f_ext=jnp.atleast_2d(forces.squeeze()).astype(float)
-                    )
-                ),
+                _link_forces=jnp.atleast_2d(forces.squeeze()).astype(float),
             )
 
         # In this case, we allow only to set the inertial 6D forces to all links
@@ -389,11 +400,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
             W_f_L = f_L
 
-            W_f0_L = (
-                jnp.zeros_like(W_f_L)
-                if not additive
-                else self.input.physics_model.f_ext
-            )
+            W_f0_L = jnp.zeros_like(W_f_L) if not additive else self._link_forces
 
             return replace(forces=W_f0_L + W_f_L)
 
@@ -410,18 +417,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Compute the bias depending on whether we either set or add the link forces.
         W_f0_L = (
-            jnp.zeros_like(f_L)
-            if not additive
-            else self.input.physics_model.f_ext[link_idxs, :]
+            jnp.zeros_like(f_L) if not additive else self._link_forces[link_idxs, :]
         )
 
         # If inertial-fixed representation, we can directly store the link forces.
         if self.velocity_representation is VelRepr.Inertial:
             W_f_L = f_L
             return replace(
-                forces=self.input.physics_model.f_ext.at[link_idxs, :].set(
-                    W_f0_L + W_f_L
-                )
+                forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L)
             )
 
         if data is None:
@@ -450,9 +453,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         W_H_L = js.model.forward_kinematics(model=model, data=data)
         W_f_L = convert_using_link_frame(f_L=f_L, W_H_L=W_H_L[link_idxs, :, :])
 
-        return replace(
-            forces=self.input.physics_model.f_ext.at[link_idxs, :].set(W_f0_L + W_f_L)
-        )
+        return replace(forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L))
 
     def apply_frame_forces(
         self,

jaxsim/rbda/contacts/visco_elastic.py

@@ -851,7 +851,7 @@ class ViscoElasticContacts(common.ContactModel):
         W_f̅_L = (
             jnp.array(average_link_contact_forces_inertial)
             if average_link_contact_forces_inertial is not None
-            else jnp.zeros_like(references.input.physics_model.f_ext)
+            else jnp.zeros_like(references._link_forces)
         ).astype(float)
 
         LW_f̿_L = (

{jaxsim-0.5.1.dev48.dist-info → jaxsim-0.5.1.dev60.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev48
+Version: 0.5.1.dev60
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev48.dist-info → jaxsim-0.5.1.dev60.dist-info}/RECORD

@@ -1,21 +1,21 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=4Uq3N7DcVq205gBgiBZ4x-nQK8MM7dxohj0UxVetWwQ,426
+jaxsim/_version.py,sha256=SbvsohyIl1wURDKT1I6oE7EjNu_dZVzu8-rPXwguX6s,426
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
-jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
-jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
-jaxsim/api/contact.py,sha256=P62D4MpgcEJScmxuhh99YcoW4QLU6hVjG_lvlLL-mxg,25214
-jaxsim/api/data.py,sha256=pGgQ8k_QJ0kh7igMmQGfccix3wOXKcSfsLnZQyOQHjo,29860
-jaxsim/api/frame.py,sha256=MT6Qpg5uJDnhik899QNpHskZMijgjMSOw-2v-t-kul4,14434
-jaxsim/api/joint.py,sha256=8rCIxRMeAidsaBbw7kkGp6z3-UmBPtqmYmV_arHDQJ8,7365
+jaxsim/api/com.py,sha256=eZEF2nAwe8FdEcTz9v--VUn6mTSRjI5nGtRpeXXW_AA,13826
+jaxsim/api/common.py,sha256=SvEOGxCKOxKLLVHaNp1sFkBX0sku3-wH0-HUlYVWCDk,7090
+jaxsim/api/contact.py,sha256=GFPLolO4feVUSWrLbHwvV64bD14ESzGFhoiCTJ_Im2Q,25421
+jaxsim/api/data.py,sha256=hz8g-P0o7XoDYyUFPx6yA8QlgHmjfFf2_OdiwcRV6W8,30292
+jaxsim/api/frame.py,sha256=7V6ih6cUe5cBKp8GKD0mZ7Zm-e3lj7qiffoAcVgZ76o,14595
+jaxsim/api/joint.py,sha256=P4j_5caFs_YkV_jHI_98hko0Q5-CmWAYCqHgr8eGILY,7457
 jaxsim/api/kin_dyn_parameters.py,sha256=wnto0nzzEJ_M8tH2PUdldEyxQwQdsStYUoQFu696uuw,29897
 jaxsim/api/link.py,sha256=nHjffhNdi_xGkteMsqdb_hC9mdV9rNw7k3pl89Uhw_8,12798
-jaxsim/api/model.py,sha256=0Ve1zEVoe5mFw4kX5sG6eCg0PXRnOZl1f-wsmFpAxeM,79659
-jaxsim/api/ode.py,sha256=_t18avoCJngQk6eMFTGpaeahbpchQP20qJnUOCPkz8s,15360
-jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
-jaxsim/api/references.py,sha256=eIOk3MAOc9LJSKfI8M4WA8gGD-meo50vRfhXdea4sNI,20539
+jaxsim/api/model.py,sha256=ccS7oIpE6LKFEVG4dvXktD_zGQOoieKULANGjY8TZUs,80139
+jaxsim/api/ode.py,sha256=3pSKJaZ6NrtADuulUUNmp_828PDXELgENPjhQ9UeJCc,15429
+jaxsim/api/ode_data.py,sha256=ggF1AVaLW5QuXrfpNsFs-voVcW6gZkxK2Xe9GiDmou0,13755
+jaxsim/api/references.py,sha256=YkdZhRv8NoBC94qvpwn1w9_alVuxrfiZV5w5NHQIt-g,20737
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=ohISUnUWTaNHt2kweg1JyzwYGZgIH_wc-01qJWJsO80,18281
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
@@ -59,15 +59,15 @@ jaxsim/rbda/contacts/common.py,sha256=BjwZMCkzd1ZOdZW7_Zt09Cl5j2JUHXM5Q8ao_qS6e6
 jaxsim/rbda/contacts/relaxed_rigid.py,sha256=PgwKfProN5sLXJsSov3nIidHHMVpJqIp7eIv6_bPGjs,20345
 jaxsim/rbda/contacts/rigid.py,sha256=X-PE6PmZqlKoZTY6JhYBSW-vom-rq2uBKmBUNQeQHCg,15991
 jaxsim/rbda/contacts/soft.py,sha256=sIWT4NUJmoVR5T1Fo0ExdPfzL_gPfiPiB-9CFuotE_s,15567
-jaxsim/rbda/contacts/visco_elastic.py,sha256=4gVs3oVzn1IImAsnC7yZ3b-aMwaA-7bSbYqxTULY1FE,39899
+jaxsim/rbda/contacts/visco_elastic.py,sha256=QhyJHjDowyBTAhoSdZcCIkOqzp__gMXhLON-qYyMgQc,39886
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=_G1QS3zWycj089R8fTP5s2VjcZpEdJxREjXZJ-oXIvc,5248
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=eEY28MZW0Lm_jJNt1NkFqZz0ek01tvhR46OXZYCo7tc,532
 jaxsim/utils/wrappers.py,sha256=ZY7olSORzZRvSzkdeNLj8yjwUIAt9L0Douwl7wItjpk,4008
-jaxsim-0.5.1.dev48.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.5.1.dev48.dist-info/METADATA,sha256=OgcJvWF1ohXuO7dpnmjcY4upeZhqyGwMPcTGkvfAOF4,17937
-jaxsim-0.5.1.dev48.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-jaxsim-0.5.1.dev48.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.5.1.dev48.dist-info/RECORD,,
+jaxsim-0.5.1.dev60.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.5.1.dev60.dist-info/METADATA,sha256=_jyXanq6vSSkWmVG8N8_5WrIPq-VQx4Skaq4xb-lvFw,17937
+jaxsim-0.5.1.dev60.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+jaxsim-0.5.1.dev60.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.5.1.dev60.dist-info/RECORD,,