jaxsim 0.1rc0__py3-none-any.whl → 0.2.dev8__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1rc0'
-__version_tuple__ = version_tuple = (0, 1)
+__version__ = version = '0.2.dev8'
+__version_tuple__ = version_tuple = (0, 2, 'dev8')

jaxsim/high_level/model.py

@@ -385,13 +385,13 @@ class Model(Vmappable):
     def link_names(self) -> tuple[str, ...]:
         """"""
 
-        return tuple(l.name() for l in self.links())
+        return tuple(self.physics_model.description.links_dict.keys())
 
     @functools.partial(oop.jax_tf.method_ro, jit=False, vmap=False)
     def joint_names(self) -> tuple[str, ...]:
         """"""
 
-        return tuple(j.name() for j in self.joints())
+        return tuple(self.physics_model.description.joints_dict.keys())
 
     @functools.partial(oop.jax_tf.method_ro, jit=False, vmap=False)
     def links(

jaxsim/parsers/rod/parser.py

@@ -135,11 +135,13 @@ def extract_model_data(
                 parent=world_link,
                 child=links_dict[j.child],
                 jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-                axis=np.array(j.axis.xyz.xyz)
-                if j.axis is not None
-                and j.axis.xyz is not None
-                and j.axis.xyz.xyz is not None
-                else None,
+                axis=(
+                    np.array(j.axis.xyz.xyz)
+                    if j.axis is not None
+                    and j.axis.xyz is not None
+                    and j.axis.xyz.xyz is not None
+                    else None
+                ),
                 pose=j.pose.transform() if j.pose is not None else np.eye(4),
             )
             for j in sdf_model.joints()
@@ -200,41 +202,55 @@ def extract_model_data(
             parent=links_dict[j.parent],
             child=links_dict[j.child],
             jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-            axis=np.array(j.axis.xyz.xyz)
-            if j.axis is not None
-            and j.axis.xyz is not None
-            and j.axis.xyz.xyz is not None
-            else None,
+            axis=(
+                np.array(j.axis.xyz.xyz)
+                if j.axis is not None
+                and j.axis.xyz is not None
+                and j.axis.xyz.xyz is not None
+                else None
+            ),
             pose=j.pose.transform() if j.pose is not None else np.eye(4),
             initial_position=0.0,
             position_limit=(
-                float(j.axis.limit.lower)
-                if j.axis is not None and j.axis.limit is not None
-                else np.finfo(float).min,
-                float(j.axis.limit.upper)
-                if j.axis is not None and j.axis.limit is not None
-                else np.finfo(float).max,
+                (
+                    float(j.axis.limit.lower)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).min
+                ),
+                (
+                    float(j.axis.limit.upper)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).max
+                ),
+            ),
+            friction_static=(
+                j.axis.dynamics.friction
+                if j.axis is not None
+                and j.axis.dynamics is not None
+                and j.axis.dynamics.friction is not None
+                else 0.0
+            ),
+            friction_viscous=(
+                j.axis.dynamics.damping
+                if j.axis is not None
+                and j.axis.dynamics is not None
+                and j.axis.dynamics.damping is not None
+                else 0.0
+            ),
+            position_limit_damper=(
+                j.axis.limit.dissipation
+                if j.axis is not None
+                and j.axis.limit is not None
+                and j.axis.limit.dissipation is not None
+                else 0.0
+            ),
+            position_limit_spring=(
+                j.axis.limit.stiffness
+                if j.axis is not None
+                and j.axis.limit is not None
+                and j.axis.limit.stiffness is not None
+                else 0.0
             ),
-            friction_static=j.axis.dynamics.friction
-            if j.axis is not None
-            and j.axis.dynamics is not None
-            and j.axis.dynamics.friction is not None
-            else 0.0,
-            friction_viscous=j.axis.dynamics.damping
-            if j.axis is not None
-            and j.axis.dynamics is not None
-            and j.axis.dynamics.damping is not None
-            else 0.0,
-            position_limit_damper=j.axis.limit.dissipation
-            if j.axis is not None
-            and j.axis.limit is not None
-            and j.axis.limit.dissipation is not None
-            else 0.0,
-            position_limit_spring=j.axis.limit.stiffness
-            if j.axis is not None
-            and j.axis.limit is not None
-            and j.axis.limit.stiffness is not None
-            else 0.0,
         )
         for j in sdf_model.joints()
         if j.type in {"revolute", "prismatic", "fixed"}

jaxsim/physics/model/physics_model.py

@@ -45,14 +45,14 @@ class PhysicsModel(JaxsimDataclass):
     )
     is_floating_base: Static[bool] = dataclasses.field(default=False)
     gc: GroundContact = dataclasses.field(default_factory=lambda: GroundContact())
-    description: Static[
-        jaxsim.parsers.descriptions.model.ModelDescription
-    ] = dataclasses.field(default=None)
+    description: Static[jaxsim.parsers.descriptions.model.ModelDescription] = (
+        dataclasses.field(default=None)
+    )
 
     _parent_array_dict: Static[Dict[int, int]] = dataclasses.field(default_factory=dict)
-    _jtype_dict: Static[
-        Dict[int, Union[JointType, JointDescriptor]]
-    ] = dataclasses.field(default_factory=dict)
+    _jtype_dict: Static[Dict[int, Union[JointType, JointDescriptor]]] = (
+        dataclasses.field(default_factory=dict)
+    )
     _tree_transforms_dict: Dict[int, jtp.Matrix] = dataclasses.field(
         default_factory=dict
     )

jaxsim/simulation/simulator.py

@@ -432,8 +432,9 @@ class JaxSim(Vmappable):
     def step_over_horizon(
         self,
         horizon_steps: jtp.Int,
-        callback_handler: Union["scb.SimulatorCallback", "scb.CallbackHandler"]
-        | None = None,
+        callback_handler: (
+            Union["scb.SimulatorCallback", "scb.CallbackHandler"] | None
+        ) = None,
         clear_inputs: jtp.Bool = False,
     ) -> Union[
         "JaxSim",
@@ -459,10 +460,8 @@ class JaxSim(Vmappable):
         sim = self.copy().mutable(validate=True)
 
         # Helper to get callbacks from the handler
-        get_cb = (
-            lambda h, cb_name: getattr(h, cb_name)
-            if h is not None and hasattr(h, cb_name)
-            else None
+        get_cb = lambda h, cb_name: (
+            getattr(h, cb_name) if h is not None and hasattr(h, cb_name) else None
         )
 
         # Get the callbacks

jaxsim/utils/oop.py

@@ -3,16 +3,20 @@ import dataclasses
 import functools
 import inspect
 import os
-from typing import Any, Callable, Generator
+from typing import Any, Callable, Generator, TypeVar
 
 import jax
 import jax.flatten_util
+from typing_extensions import ParamSpec
 
 from jaxsim import logging
 from jaxsim.utils import tracing
 
 from . import Mutability, Vmappable
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
 
 class jax_tf:
     """
@@ -27,13 +31,13 @@ class jax_tf:
 
     @staticmethod
     def method_ro(
-        fn: Callable,
+        fn: Callable[_P, _R],
         jit: bool = True,
         static_argnames: tuple[str, ...] | list[str] = (),
         vmap: bool | None = None,
         vmap_in_axes: tuple[int, ...] | int | None = None,
         vmap_out_axes: tuple[int, ...] | int | None = None,
-    ):
+    ) -> Callable[_P, _R]:
         """
         Decorator for r/o methods of classes inheriting from Vmappable.
         """
@@ -51,14 +55,14 @@ class jax_tf:
 
     @staticmethod
     def method_rw(
-        fn: Callable,
+        fn: Callable[_P, _R],
         validate: bool = True,
         jit: bool = True,
         static_argnames: tuple[str, ...] | list[str] = (),
         vmap: bool | None = None,
         vmap_in_axes: tuple[int, ...] | int | None = None,
         vmap_out_axes: tuple[int, ...] | int | None = None,
-    ):
+    ) -> Callable[_P, _R]:
         """
         Decorator for r/w methods of classes inheriting from Vmappable.
         """
@@ -76,7 +80,7 @@ class jax_tf:
 
     @staticmethod
     def method(
-        fn: Callable,
+        fn: Callable[_P, _R],
         read_only: bool = True,
         validate: bool = True,
         jit_enabled: bool = True,
@@ -109,7 +113,7 @@ class jax_tf:
         """
 
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs):
             """The wrapper function that is returned by this decorator."""
 
             # Methods of classes inheriting from Vmappable decorated by this wrapper
@@ -202,9 +206,9 @@ class jax_tf:
             mutability_dict = {
                 Mutability.MUTABLE_NO_VALIDATION: Mutability.MUTABLE_NO_VALIDATION,
                 Mutability.MUTABLE: Mutability.MUTABLE,
-                Mutability.FROZEN: Mutability.MUTABLE
-                if validate
-                else Mutability.MUTABLE_NO_VALIDATION,
+                Mutability.FROZEN: (
+                    Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
+                ),
             }
 
             # We need to replace all the dynamic leafs of the original instance with those

{jaxsim-0.1rc0.dist-info → jaxsim-0.2.dev8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.1rc0
+Version: 0.2.dev8
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
@@ -38,7 +38,7 @@ Requires-Dist: jax-dataclasses >=1.4.0
 Requires-Dist: pptree
 Requires-Dist: rod
 Provides-Extra: all
-Requires-Dist: black ; extra == 'all'
+Requires-Dist: black[jupyter] ; extra == 'all'
 Requires-Dist: isort ; extra == 'all'
 Requires-Dist: idyntree ; extra == 'all'
 Requires-Dist: pytest >=6.0 ; extra == 'all'
@@ -46,7 +46,7 @@ Requires-Dist: pytest-forked ; extra == 'all'
 Requires-Dist: pytest-icdiff ; extra == 'all'
 Requires-Dist: robot-descriptions ; extra == 'all'
 Provides-Extra: style
-Requires-Dist: black ; extra == 'style'
+Requires-Dist: black[jupyter] ; extra == 'style'
 Requires-Dist: isort ; extra == 'style'
 Provides-Extra: testing
 Requires-Dist: idyntree ; extra == 'testing'

{jaxsim-0.1rc0.dist-info → jaxsim-0.2.dev8.dist-info}/RECORD

@@ -1,12 +1,12 @@
 jaxsim/__init__.py,sha256=LJhCG4rsmCrTKTocwRIvllPQeYTxDn-VFn6NjPngn4s,1877
-jaxsim/_version.py,sha256=nrBgoAeJilYeqm-4jM5kvaCqeYEJ_yzl6oqlUk8BZK0,409
+jaxsim/_version.py,sha256=PIvMaj0I1p6HezVGBy1PEqXZf5zmL-w3B5rbYTFjdUY,419
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=Skdm3OrTCT0MvaL57kD10-6LAfceYIUTzO8zuP5b0RA,777
 jaxsim/high_level/__init__.py,sha256=aWYBCsYmEO76Qt4GEi91Hye_ifGFLvc_bpy9OQplz2o,69
 jaxsim/high_level/common.py,sha256=6nyRlFsNOLEy5JvLH70VPWeHGSL_ZKNxX3Q62ccqSuY,196
 jaxsim/high_level/joint.py,sha256=0WF0QWkZzP0SXw0QYpn3PAwdZq0_uXFr2_f1OATiOBA,4089
 jaxsim/high_level/link.py,sha256=4kcBMh-3w9c-fkTYm3_sXfdwd3NwUm7jKf5BjwEge94,8010
-jaxsim/high_level/model.py,sha256=Krv1ZPWWz7bL_XdeyupHXoGLDI0fZy7aQcJPKpJJS2w,57347
+jaxsim/high_level/model.py,sha256=oyLMHkhIOqXL_cnj98wG_npHCsBhiLd5uF3AFcrTZqw,57383
 jaxsim/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/math/adjoint.py,sha256=ImkOkWHQKMukBprLTsOPpSuqb1NNPA3_t447zRVo79s,3779
 jaxsim/math/conv.py,sha256=jbr9MU_vGtTpLTQpbqwAT4huF51skCfORFUBmKSdlaI,3138
@@ -25,7 +25,7 @@ jaxsim/parsers/descriptions/joint.py,sha256=hpH0ANvIhbEQk-NGRmWIvPv3lXW385TBIMWN
 jaxsim/parsers/descriptions/link.py,sha256=SONPaSwtNhIX93RIVng8Fb_Y7I5h3sk-5rsqd5U7Fmw,2493
 jaxsim/parsers/descriptions/model.py,sha256=wenuDrjoBf6prkzm9WyYT0nFWc0l6WBpKNjLoRUDPxo,8937
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=JcHkCoIhsecA1MVfn8Q0Ubkhyzb5FDpaiZzFh-rbQ6Q,12173
+jaxsim/parsers/rod/parser.py,sha256=_lbU3Oaj7-2GDw-KxCCjWB4ZpAj0I5oEDxquN5Ff0p4,12565
 jaxsim/parsers/rod/utils.py,sha256=C3RfLSnHHR7rgZhnM15QjM_tpJFHKf1Jf2mOsdep3LM,6492
 jaxsim/physics/__init__.py,sha256=TKH7IqQi39eobWT03b820ky6tWVWFFWcO8YPayNpnZc,216
 jaxsim/physics/algos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -41,24 +41,24 @@ jaxsim/physics/algos/terrain.py,sha256=Gw9-1AjU4c4Yd2xzo0i-fgWwYlroj03TjScJsz_2m
 jaxsim/physics/algos/utils.py,sha256=0OiELbXtv5Jink3l-vMK_OGHgGkZ_wTAAclcd7vDKoc,2230
 jaxsim/physics/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/physics/model/ground_contact.py,sha256=mva-yDzYHREmgUu8jGJmIAsf66_SF6ZISmN-XQQ9Ktw,1924
-jaxsim/physics/model/physics_model.py,sha256=UtGBTxXR0cLEc9EHzDo6Q__RRNgX4VCS88wQ7SsJuyU,13245
+jaxsim/physics/model/physics_model.py,sha256=kVTIaJQrxALzyWjWrDLnwDOcxmzaPGSpUOS8BCq-g6M,13249
 jaxsim/physics/model/physics_model_state.py,sha256=nAXAwXnL7kez8DlxI_AxOY-p9ZJjVFAgZ-P0JZ5-crA,5384
 jaxsim/simulation/__init__.py,sha256=WOWkzq7rMGa4xWvjNqTYtD0Nl4yLQtULGW1xU7hD9m0,182
 jaxsim/simulation/integrators.py,sha256=E6kz4irJYzPi4hUnzyxOgGZJiOb1ATT9QQ-UtyLOjMo,13779
 jaxsim/simulation/ode.py,sha256=cawB4qsuSoGifF_wkwGwOFU8NB0R89isa7sK-IZipv8,10376
 jaxsim/simulation/ode_data.py,sha256=spzHU5LnOL6mJPuuhho-J61koT-bcTRonqMMkiPo3M4,1750
 jaxsim/simulation/ode_integration.py,sha256=56PqhI_MEeiBQ3MNpE2ytSQBntBPPlUM-oe50DOpbHw,4042
-jaxsim/simulation/simulator.py,sha256=Et58GPQkXIWL0NfM8DczbjJCixz2KI6Imz8w9KFWzN4,18168
+jaxsim/simulation/simulator.py,sha256=2wy6cVmwgMk6AVZ8GJ8jdLIfCRMzLe64wwqUCMtL9Xs,18160
 jaxsim/simulation/simulator_callbacks.py,sha256=FEKyQyhVmDRckIOizvNDhSj1Xh6gYLO7yXDrK09_VQY,1476
 jaxsim/simulation/utils.py,sha256=YdNA1mYGBAE7xVA-Dw7_OoBEuh0J8RS2X0RPQZf4c5E,329
 jaxsim/sixd/__init__.py,sha256=3tbynXQjvJ6X1IRcDH5eQBgBL0ilTSerDkS8SEF7a8A,62
 jaxsim/utils/__init__.py,sha256=UQsjrWMrhSQUfUXIIQhzd0kEioLHR3U0OhB-sIQqOd4,291
 jaxsim/utils/jaxsim_dataclass.py,sha256=FbjfEoCoYC_F-M3wUggXiEhQ7MMS-V_ciYQca-uSiMQ,3272
-jaxsim/utils/oop.py,sha256=PPQvD2_pXx81UD4dbu8E8RiYLegJNEsD23IT-YkEuQs,22434
+jaxsim/utils/oop.py,sha256=LQhBXkSOD0zgYNJLO7Bl0FPRg-LvtvPzxyQa1WFP0rM,22616
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/vmappable.py,sha256=NqGL9nGFRI5OorCfnjXsjR_yXigzDxL0lW1YhQ_nMTY,3655
-jaxsim-0.1rc0.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
-jaxsim-0.1rc0.dist-info/METADATA,sha256=lCoeEqU6jjUIR1o02r7qhR4akKSFw9h-Q56uFT-12gM,7143
-jaxsim-0.1rc0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-jaxsim-0.1rc0.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.1rc0.dist-info/RECORD,,
+jaxsim-0.2.dev8.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
+jaxsim-0.2.dev8.dist-info/METADATA,sha256=8OFNTQyRWt1A5WWP4-K6qtGGFTU-xNvnC34W-0emyys,7163
+jaxsim-0.2.dev8.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+jaxsim-0.2.dev8.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.dev8.dist-info/RECORD,,