jaxsim 0.3.1.dev48__py3-none-any.whl → 0.3.1.dev51__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/model.py

@@ -1841,7 +1841,13 @@ def step(
         and the new state of the integrator.
     """
 
-    integrator_kwargs = kwargs if kwargs is not None else dict()
+    # Extract the integrator kwargs.
+    # The following logic allows using integrators having kwargs colliding with the
+    # kwargs of this step function.
+    kwargs = kwargs if kwargs is not None else {}
+    integrator_kwargs = kwargs.pop("integrator_kwargs", {})
+    integrator_kwargs = kwargs | integrator_kwargs
+
     integrator_state = integrator_state if integrator_state is not None else dict()
 
     # Extract the initial resources.
@@ -1855,8 +1861,21 @@ def step(
         t0=jnp.array(t0_ns / 1e9).astype(float),
         dt=dt,
         params=integrator_state_x0,
+        # Always inject the current (model, data) pair into the system dynamics
+        # considered by the integrator, and include the input variables represented
+        # by the pair (joint_forces, link_forces).
+        # Note that the wrapper of the system dynamics will override (state_x0, t0)
+        # inside the passed data even if it is not strictly needed. This logic is
+        # necessary to re-use the jit-compiled step function of compatible pytrees
+        # of model and data produced e.g. by parameterized applications.
         **(
-            dict(joint_forces=joint_forces, link_forces=link_forces) | integrator_kwargs
+            dict(
+                model=model,
+                data=data,
+                joint_forces=joint_forces,
+                link_forces=link_forces,
+            )
+            | integrator_kwargs
         ),
     )
 

jaxsim/integrators/common.py

@@ -422,7 +422,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_map(lambda l: l[self.index_of_fsal], K)
+            self.params["dxdt0"] = jax.tree_util.tree_map(
+                lambda l: l[self.index_of_fsal], K
+            )
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.

{jaxsim-0.3.1.dev48.dist-info → jaxsim-0.3.1.dev51.dist-info}/METADATA

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

{jaxsim-0.3.1.dev48.dist-info → jaxsim-0.3.1.dev51.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=xzuTuZrgKdWLqqDzbvqzm2cJrEtAbepOeUqDu7ByVek,2621
-jaxsim/_version.py,sha256=1H8tF2jqYwy0-hfcKnbgbUmsunSq6OC6nWm1GaH-kTQ,426
+jaxsim/_version.py,sha256=IGT0sukgd0adIU8Kj5QFc30r0PeT_NCk5YVNZH5UbRg,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=cl7HHQCeP3mHmtF6EuQZcCjGvDmc_AryMWntP_lRBGg,722
@@ -12,12 +12,12 @@ jaxsim/api/frame.py,sha256=m_waB9_0kgJq5miiZDXdRzIZii-BwQaN9bRt22JkJ1I,7212
 jaxsim/api/joint.py,sha256=Pvg_It2iYA-jAQ2nOlFZxwmITiozO_f46G13BdQtHQ0,5106
 jaxsim/api/kin_dyn_parameters.py,sha256=AEpDg9kihbKUN9PA8pNrAruSuWFUC-k_GGxtlcdcDiQ,29215
 jaxsim/api/link.py,sha256=edXaNO0TcqyFyMOlIlCnReQK_VP8p38crLEp0of7mWo,18404
-jaxsim/api/model.py,sha256=HAnrlgPDl5CCZQzQ84AfjC_DZjmrCzBKEDodE6hyLf8,60518
+jaxsim/api/model.py,sha256=G4tlJ0mkpH_v31ZPObBSeTqaioYJpdfapmIG7-FyWe8,61424
 jaxsim/api/ode.py,sha256=xQL53ppnKweMQWRNm5gGR8FTjqRVzds8WKg9js9k5TA,10780
 jaxsim/api/ode_data.py,sha256=Sa2i1zZhqyQqIGv1jarTmmU-W9HhTw-DErs12kFA1GA,19737
 jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=9HXRVFo95Mpt6RcVhBrOfvOO7mDxqbkXeg_lKUibEFY,20693
+jaxsim/integrators/common.py,sha256=cRXD0UPEFo2UkPQnOUi3u9Ph3kbfaYFaAvGnZWjSWwI,20733
 jaxsim/integrators/fixed_step.py,sha256=JXaEyEzfSiYea0GnPA7l27J3X0YPB0e25D4qfrxAvzQ,2766
 jaxsim/integrators/variable_step.py,sha256=jq3PStzFiMciU7lux6CTj4B3gVOfSpYgK2oz2yzIbdo,21380
 jaxsim/math/__init__.py,sha256=inJ9nRFkqstuGa8OyFkfWVudo5U9Ug4WgDBuKva8AIA,337
@@ -61,8 +61,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=QIJitSoljrKR_U4T3ewCJPT3DTh-tPZsRsg0t_MH93E,3896
-jaxsim-0.3.1.dev48.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.3.1.dev48.dist-info/METADATA,sha256=Dn56OLcqYLKMeb2dcuR6B8BD_Lq_gcb0rE5QwfJ1TS0,9739
-jaxsim-0.3.1.dev48.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.3.1.dev48.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.3.1.dev48.dist-info/RECORD,,
+jaxsim-0.3.1.dev51.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.3.1.dev51.dist-info/METADATA,sha256=zbBiTRoZ6XIJvJcq8UTkHhvGOc7krPxFpxDqz5_oJ_A,9739
+jaxsim-0.3.1.dev51.dist-info/WHEEL,sha256=cpQTJ5IWu9CdaPViMhC9YzF8gZuS5-vlfoFihTBC86A,91
+jaxsim-0.3.1.dev51.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.3.1.dev51.dist-info/RECORD,,

{jaxsim-0.3.1.dev48.dist-info → jaxsim-0.3.1.dev51.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (70.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any