jaxsim 0.4.2.dev45__py3-none-any.whl → 0.4.2.dev50__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.2.dev45'
-__version_tuple__ = version_tuple = (0, 4, 2, 'dev45')
+__version__ = version = '0.4.2.dev50'
+__version_tuple__ = version_tuple = (0, 4, 2, 'dev50')

jaxsim/api/references.py

@@ -9,7 +9,6 @@ import jax_dataclasses
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim import exceptions
-from jaxsim.math import Adjoint
 from jaxsim.utils.tracing import not_tracing
 
 from .common import VelRepr
@@ -493,9 +492,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Extract the frame indices.
         frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
-        parent_link_idxs = jax.vmap(
-            lambda frame_idx: js.frame.idx_of_parent_link, in_axes=(None,)
-        )(model, frame_idx=frame_idxs)
+        parent_link_idxs = jax.vmap(js.frame.idx_of_parent_link, in_axes=(None,))(
+            model, frame_index=frame_idxs
+        )
 
         exceptions.raise_value_error_if(
             condition=jnp.logical_not(data.valid(model=model)),
@@ -527,25 +526,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             case _:
                 raise ValueError("Invalid velocity representation.")
 
-        W_H_L = js.model.forward_kinematics(model=model, data=data)
-
-        def convert_to_link_force(
-            W_f_F: jtp.MatrixLike, W_H_F: jtp.MatrixLike, parent_link_idx: jtp.ArrayLike
-        ) -> jtp.Matrix:
-            L_Xf_W = Adjoint.from_transform(W_H_L[parent_link_idx]).T
-
-            return L_Xf_W @ W_f_F
-
-        W_f_L_i = jax.vmap(convert_to_link_force)(W_f_F, W_H_Fi, parent_link_idxs)
-
         # Sum the forces on the parent links.
         mask = parent_link_idxs[:, jnp.newaxis] == jnp.arange(model.number_of_links())
-        W_f_L = mask.T @ W_f_L_i
+        W_f_L = mask.T @ W_f_F
 
         with self.switch_velocity_representation(
             velocity_representation=VelRepr.Inertial
         ):
-            return self.apply_link_forces(
+            references = self.apply_link_forces(
                 model=model,
                 data=data,
                 link_names=js.link.idxs_to_names(
@@ -554,3 +542,8 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
                 forces=W_f_L,
                 additive=additive,
             )
+
+        with references.switch_velocity_representation(
+            velocity_representation=self.velocity_representation
+        ):
+            return references

jaxsim/rbda/crba.py

@@ -59,10 +59,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         return (i_X_0,), None
 
-    (i_X_0,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=forward_pass_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (i_X_0,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=forward_pass_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(i_X_0,), None]
     )
 
     # ===================
@@ -128,10 +132,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
                 operand=carry,
             )
 
-        (j, Fi, M), _ = jax.lax.scan(
-            f=inner_fn,
-            init=carry_inner_fn,
-            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        (j, Fi, M), _ = (
+            jax.lax.scan(
+                f=inner_fn,
+                init=carry_inner_fn,
+                xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+            )
+            if model.number_of_links() > 1
+            else [(j, Fi, M), None]
         )
 
         Fi = i_X_0[j].T @ Fi
@@ -143,10 +151,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
     # This scan performs the backward pass to compute Mbj, Mjb and Mjj, that
     # also includes a fake while loop implemented with a scan and two cond.
-    (Mc, M), _ = jax.lax.scan(
-        f=backward_pass,
-        init=backward_pass_carry,
-        xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+    (Mc, M), _ = (
+        jax.lax.scan(
+            f=backward_pass,
+            init=backward_pass_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(Mc, M), None]
     )
 
     # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶.

jaxsim/rbda/forward_kinematics.py

@@ -75,10 +75,14 @@ def forward_kinematics_model(
 
         return (W_X_i,), None
 
-    (W_X_i,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (W_X_i,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_kinematics_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i,), None]
     )
 
     return jax.vmap(Adjoint.to_transform)(W_X_i)

jaxsim/rbda/jacobian.py

@@ -67,10 +67,14 @@ def jacobian(
 
         return (i_X_0,), None
 
-    (i_X_0,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    (i_X_0,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_kinematics_carry,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(i_X_0,), None]
     )
 
     # ============================
@@ -105,10 +109,14 @@ def jacobian(
 
         return J, None
 
-    L_J_WL_B, _ = jax.lax.scan(
-        f=compute_jacobian,
-        init=J,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    L_J_WL_B, _ = (
+        jax.lax.scan(
+            f=compute_jacobian,
+            init=J,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [J, None]
     )
 
     return L_J_WL_B
@@ -184,10 +192,14 @@ def jacobian_full_doubly_left(
 
         return (B_X_i, J), None
 
-    (B_X_i, J), _ = jax.lax.scan(
-        f=compute_full_jacobian,
-        init=compute_full_jacobian_carry,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    (B_X_i, J), _ = (
+        jax.lax.scan(
+            f=compute_full_jacobian,
+            init=compute_full_jacobian_carry,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(B_X_i, J), None]
     )
 
     # Convert adjoints to SE(3) transforms.

{jaxsim-0.4.2.dev45.dist-info → jaxsim-0.4.2.dev50.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev45
+Version: 0.4.2.dev50
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>
@@ -60,29 +60,29 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
-Requires-Dist: jax >=0.4.13
-Requires-Dist: jaxlib >=0.4.13
-Requires-Dist: jaxlie >=1.3.0
-Requires-Dist: jax-dataclasses >=1.4.0
+Requires-Dist: jax>=0.4.13
+Requires-Dist: jaxlib>=0.4.13
+Requires-Dist: jaxlie>=1.3.0
+Requires-Dist: jax-dataclasses>=1.4.0
 Requires-Dist: pptree
-Requires-Dist: rod >=0.3.0
-Requires-Dist: typing-extensions ; python_version < "3.12"
+Requires-Dist: rod>=0.3.0
+Requires-Dist: typing-extensions; python_version < "3.12"
 Provides-Extra: all
-Requires-Dist: jaxsim[style,testing,viz] ; extra == 'all'
+Requires-Dist: jaxsim[style,testing,viz]; extra == "all"
 Provides-Extra: style
-Requires-Dist: black[jupyter] ~=24.0 ; extra == 'style'
-Requires-Dist: isort ; extra == 'style'
-Requires-Dist: pre-commit ; extra == 'style'
+Requires-Dist: black[jupyter]~=24.0; extra == "style"
+Requires-Dist: isort; extra == "style"
+Requires-Dist: pre-commit; extra == "style"
 Provides-Extra: testing
-Requires-Dist: idyntree >=12.2.1 ; extra == 'testing'
-Requires-Dist: pytest >=6.0 ; extra == 'testing'
-Requires-Dist: pytest-icdiff ; extra == 'testing'
-Requires-Dist: robot-descriptions ; extra == 'testing'
+Requires-Dist: idyntree>=12.2.1; extra == "testing"
+Requires-Dist: pytest>=6.0; extra == "testing"
+Requires-Dist: pytest-icdiff; extra == "testing"
+Requires-Dist: robot-descriptions; extra == "testing"
 Provides-Extra: viz
-Requires-Dist: lxml ; extra == 'viz'
-Requires-Dist: mediapy ; extra == 'viz'
-Requires-Dist: mujoco >=3.0.0 ; extra == 'viz'
-Requires-Dist: scipy >=1.14.0 ; extra == 'viz'
+Requires-Dist: lxml; extra == "viz"
+Requires-Dist: mediapy; extra == "viz"
+Requires-Dist: mujoco>=3.0.0; extra == "viz"
+Requires-Dist: scipy>=1.14.0; extra == "viz"
 
 # JaxSim
 

{jaxsim-0.4.2.dev45.dist-info → jaxsim-0.4.2.dev50.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=mvgmTuGOMpNiwe4eR061Vc1ZiYsIU9hMLQKJGhaZkFI,426
+jaxsim/_version.py,sha256=yzghig35pfSMhh0vIsIOBH-eqIGj9fpA6g-9yTfJMlQ,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
@@ -15,7 +15,7 @@ jaxsim/api/link.py,sha256=GlnY7LMne-siFyg9J49IZGhiPQzS9Uk6rzQ0jI8cD_E,18622
 jaxsim/api/model.py,sha256=EdSjpKXd4N72wYjg5o0wGKFxjVMyrXg6LnlPEi3JqnU,63094
 jaxsim/api/ode.py,sha256=NnLTBvpaT4kXnbjAghXIzLv9DTMJ8bele2iOlUQDv3Q,11028
 jaxsim/api/ode_data.py,sha256=9YZX-SK_KJtoIqG-zYWZsQInb2NA_LtxDn-jtLqm_3U,19759
-jaxsim/api/references.py,sha256=6DG1V7dYLf088vSOOsvYSY-KNkzirD0zDJGohJbdLqc,21060
+jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=GqiyKTrAozuR6RuvVWdPF7locZQAXSEDY2AjTKpFGYM,20149
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
@@ -47,9 +47,9 @@ jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5
 jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
 jaxsim/rbda/collidable_points.py,sha256=Rmf1DhflhOTYh9mDalv0agS0CGSbmfoOybwP2KzKuJ0,4883
-jaxsim/rbda/crba.py,sha256=NhtZO48OUKKor7ddY7mB7h7a6idrmOyf0Vy4p7UCCgI,4724
-jaxsim/rbda/forward_kinematics.py,sha256=FmqhZD0hQpIuUlmpzlA-5b7EYBaLZarrhiSZ782aC3E,3357
-jaxsim/rbda/jacobian.py,sha256=I6mrlkk7Cpq3CE7k_tajOHCbT6vf2pW6vMS0TKNCnng,10725
+jaxsim/rbda/crba.py,sha256=zJSiHKRvNU98z2tT9prrWR4VU9wIZQWFwEut7mua6as,5044
+jaxsim/rbda/forward_kinematics.py,sha256=2GmEoWsrioVl_SAbKRKfhOLz57pY4aR81PKRdulqStA,3458
+jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
 jaxsim/rbda/rnea.py,sha256=LGXD6s3NigaVy4-WxoROjnbKLZcUoyFmS9UNu_4ldjo,7568
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -61,8 +61,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=fLl1tY3DDb3lpIhG6BPqA5W34hM84oFzL-5cuz8k-68,11379
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=GOJQCJc5zwzoEGZB62wnWWGvUUQlXvDxz_A2Q-hFv7c,4027
-jaxsim-0.4.2.dev45.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.2.dev45.dist-info/METADATA,sha256=gNR3PAFm00PuogCUoSG_ru548HM0k6IGQR6kSIIt2FI,17250
-jaxsim-0.4.2.dev45.dist-info/WHEEL,sha256=Rp8gFpivVLXx-k3U95ozHnQw8yDcPxmhOpn_Gx8d5nc,91
-jaxsim-0.4.2.dev45.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.2.dev45.dist-info/RECORD,,
+jaxsim-0.4.2.dev50.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.2.dev50.dist-info/METADATA,sha256=2Z9T8eA8kbRT6KsLPBUqmQLb6xvHBjLp_QI2yLMHWfk,17227
+jaxsim-0.4.2.dev50.dist-info/WHEEL,sha256=nCVcAvsfA9TDtwGwhYaRrlPhTLV9m-Ga6mdyDtuwK18,91
+jaxsim-0.4.2.dev50.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.2.dev50.dist-info/RECORD,,

{jaxsim-0.4.2.dev45.dist-info → jaxsim-0.4.2.dev50.dist-info}/WHEEL

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