jaxsim 0.4.3.dev77__py3-none-any.whl → 0.4.3.dev88__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.3.dev77'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev77')
+__version__ = version = '0.4.3.dev88'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev88')

jaxsim/api/contact.py

@@ -372,11 +372,9 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     """
 
     # Get the transforms of the parent link of all collidable points.
-    W_H_L = jax.vmap(
-        lambda parent_link_idx: js.link.transform(
-            model=model, data=data, link_index=parent_link_idx
-        )
-    )(jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int))
+    W_H_L = js.model.forward_kinematics(model=model, data=data)[
+        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
+    ]
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
@@ -427,9 +425,9 @@ def jacobian(
     # Compute the contact Jacobian.
     # In inertial-fixed output representation, the Jacobian of the parent link is also
     # the Jacobian of the frame C implicitly associated with the collidable point.
-    W_J_WC = jax.vmap(lambda parent_link_idx: W_J_WL[parent_link_idx])(
+    W_J_WC = W_J_WL[
         jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
-    )
+    ]
 
     # Adjust the output representation.
     match output_vel_repr:

jaxsim/api/kin_dyn_parameters.py

@@ -98,9 +98,7 @@ class KynDynParameters(JaxsimDataclass):
         ]
 
         # Create a vectorized object of link parameters.
-        link_parameters = jax.tree_util.tree_map(
-            lambda *l: jnp.stack(l), *link_parameters_list
-        )
+        link_parameters = jax.tree.map(lambda *l: jnp.stack(l), *link_parameters_list)
 
         # =================
         # Joints properties
@@ -114,7 +112,7 @@ class KynDynParameters(JaxsimDataclass):
 
         # Create a vectorized object of joint parameters.
         joint_parameters = (
-            jax.tree_util.tree_map(lambda *l: jnp.stack(l), *joint_parameters_list)
+            jax.tree.map(lambda *l: jnp.stack(l), *joint_parameters_list)
             if len(ordered_joints) > 0
             else JointParameters(
                 index=jnp.array([], dtype=int),
@@ -424,9 +422,7 @@ class KynDynParameters(JaxsimDataclass):
         # Note that here we include also the index 0 since suc_H_child[0] stores the
         # optional pose of the base link w.r.t. the root frame of the model.
         # This is supported by SDF when the base link <pose> element is defined.
-        suc_H_i = jax.vmap(lambda i: self.joint_model.successor_H_child(joint_index=i))(
-            jnp.arange(0, 1 + self.number_of_joints())
-        )
+        suc_H_i = self.joint_model.suc_H_i[jnp.arange(0, 1 + self.number_of_joints())]
 
         # Compute the overall transforms from the parent to the child of each joint by
         # composing all the components of our joint model.

jaxsim/api/link.py

@@ -154,7 +154,7 @@ def spatial_inertia(
         idx=link_index,
     )
 
-    link_parameters = jax.tree_util.tree_map(
+    link_parameters = jax.tree.map(
         lambda l: l[link_index], model.kin_dyn_parameters.link_parameters
     )
 

jaxsim/api/model.py

@@ -395,13 +395,7 @@ def total_mass(model: JaxSimModel) -> jtp.Float:
         The total mass of the model.
     """
 
-    return (
-        jax.vmap(lambda idx: js.link.mass(model=model, link_index=idx))(
-            jnp.arange(model.number_of_links())
-        )
-        .sum()
-        .astype(float)
-    )
+    return model.kin_dyn_parameters.link_parameters.mass.sum().astype(float)
 
 
 @jax.jit
@@ -974,7 +968,7 @@ def free_floating_coriolis_matrix(
             lambda link_index: js.link.jacobian_derivative(
                 model=model, data=data, link_index=link_index
             )
-        )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+        )(jnp.arange(model.number_of_links()))
 
     L_M_L = link_spatial_inertia_matrices(model=model)
 

jaxsim/integrators/common.py

@@ -173,9 +173,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
         # Make sure that all leafs of the dictionary are JAX arrays.
         # Also, since these are dummy parameters, set them all to zero.
-        params_after_init = jax.tree_util.tree_map(
-            lambda l: jnp.zeros_like(l), integrator.params
-        )
+        params_after_init = jax.tree.map(lambda l: jnp.zeros_like(l), integrator.params)
 
         # Mark the next step as first step after initialization.
         params_after_init = params_after_init | {
@@ -290,7 +288,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         z, aux_dict = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
 
         # The next state is the batch element located at the configured index of solution.
-        next_state = jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
+        next_state = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
 
         return next_state, aux_dict
 
@@ -327,7 +325,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         """
 
         op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
-        return jax.tree_util.tree_map(op, x0, k)
+        return jax.tree.map(op, x0, k)
 
     @classmethod
     def post_process_state(
@@ -374,7 +372,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_util.tree_map(
+        carry0 = jax.tree.map(
             lambda l: jnp.repeat(jnp.zeros_like(l)[jnp.newaxis, ...], c.size, axis=0),
             x0,
         )
@@ -398,7 +396,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
                 # Compute ∑ⱼ aᵢⱼ kⱼ.
                 op_sum_ak = lambda k: jnp.einsum("s,s...->...", A[i], k)
-                sum_ak = jax.tree_util.tree_map(op_sum_ak, K)
+                sum_ak = jax.tree.map(op_sum_ak, K)
 
                 # Compute the next state for the kᵢ evaluation.
                 # Note that this is not a Δt integration since aᵢⱼ could be fractional.
@@ -419,7 +417,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
             # Store the kᵢ derivative in K.
             op = lambda l_k, l_ki: l_k.at[i].set(l_ki)
-            K = jax.tree_util.tree_map(op, K, ki)
+            K = jax.tree.map(op, K, ki)
 
             carry = K
             return carry, aux_dict
@@ -433,14 +431,12 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_util.tree_map(
-                lambda l: l[self.index_of_fsal], K
-            )
+            self.params["dxdt0"] = jax.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`.
         op = lambda x0, k: x0 + Δt * jnp.einsum("zs,s...->z...", b.T, k)
-        z = jax.tree_util.tree_map(op, x0, K)
+        z = jax.tree.map(op, x0, K)
 
         # Transform the final state of the integration.
         # This allows to inject custom logic, if needed.

jaxsim/integrators/variable_step.py

@@ -87,13 +87,13 @@ def estimate_step_size(
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda x: atol + jnp.abs(x) * rtol
-    scale0 = jax.tree_util.tree_map(compute_scale, x0)
-    scale1 = jax.tree_util.tree_map(compute_scale, ẋ0)
+    scale0 = jax.tree.map(compute_scale, x0)
+    scale1 = jax.tree.map(compute_scale, ẋ0)
 
     # Scale the initial state and its derivative.
     scale_pytree = lambda x, scale: jnp.abs(x) / scale
-    x0_scaled = jax.tree_util.tree_map(scale_pytree, x0, scale0)
-    ẋ0_scaled = jax.tree_util.tree_map(scale_pytree, ẋ0, scale1)
+    x0_scaled = jax.tree.map(scale_pytree, x0, scale0)
+    ẋ0_scaled = jax.tree.map(scale_pytree, ẋ0, scale1)
 
     # Get the maximum of the scaled pytrees.
     d0 = jnp.linalg.norm(flatten(x0_scaled), ord=jnp.inf)
@@ -103,16 +103,16 @@ def estimate_step_size(
     h0 = jnp.where(jnp.minimum(d0, d1) <= 1e-5, 1e-6, 0.01 * d0 / d1)
 
     # Compute the next state (explicit Euler step) and its derivative.
-    x1 = jax.tree_util.tree_map(lambda x0, ẋ0: x0 + h0 * ẋ0, x0, ẋ0)
+    x1 = jax.tree.map(lambda x0, ẋ0: x0 + h0 * ẋ0, x0, ẋ0)
     ẋ1 = f(x1, t0 + h0)[0]
 
     # Compute the scaling factor of the state derivatives.
     compute_scale_2 = lambda ẋ0, ẋ1: atol + jnp.maximum(jnp.abs(ẋ0), jnp.abs(ẋ1)) * rtol
-    scale2 = jax.tree_util.tree_map(compute_scale_2, ẋ0, ẋ1)
+    scale2 = jax.tree.map(compute_scale_2, ẋ0, ẋ1)
 
     # Scale the difference of the state derivatives.
     scale_ẋ_difference = lambda ẋ0, ẋ1, scale: jnp.abs((ẋ0 - ẋ1) / scale)
-    ẋ_difference_scaled = jax.tree_util.tree_map(scale_ẋ_difference, ẋ0, ẋ1, scale2)
+    ẋ_difference_scaled = jax.tree.map(scale_ẋ_difference, ẋ0, ẋ1, scale2)
 
     # Get the maximum of the scaled derivatives difference.
     d2 = jnp.linalg.norm(flatten(ẋ_difference_scaled), ord=jnp.inf) / h0
@@ -151,11 +151,11 @@ def compute_pytree_scale(
     """
 
     # Consider a zero second pytree, if not given.
-    x2 = jax.tree_util.tree_map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
+    x2 = jax.tree.map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda l1, l2: atol + jnp.maximum(jnp.abs(l1), jnp.abs(l2)) * rtol
-    scale = jax.tree_util.tree_map(compute_scale, x1, x2)
+    scale = jax.tree.map(compute_scale, x1, x2)
 
     return scale
 
@@ -198,14 +198,14 @@ def local_error_estimation(
 
     # Consider a zero estimated final state, if not given.
     xf_estimate = (
-        jax.tree_util.tree_map(lambda l: jnp.zeros_like(l), xf)
+        jax.tree.map(lambda l: jnp.zeros_like(l), xf)
         if xf_estimate is None
         else xf_estimate
     )
 
     # Estimate the error.
     estimate_error = lambda l, l̂, sc: jnp.abs(l - l̂) / sc
-    error_estimate = jax.tree_util.tree_map(estimate_error, xf, xf_estimate, scale)
+    error_estimate = jax.tree.map(estimate_error, xf, xf_estimate, scale)
 
     # Return the highest element of the error estimate.
     return jnp.linalg.norm(flatten(error_estimate), ord=norm_ord)
@@ -359,10 +359,8 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 params_next = integrator.params
 
             # Extract the high-order solution xf and the low-order estimate x̂f.
-            xf = jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
-            x̂f = jax.tree_util.tree_map(
-                lambda l: l[self.row_index_of_solution_estimate], z
-            )
+            xf = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
+            x̂f = jax.tree.map(lambda l: l[self.row_index_of_solution_estimate], z)
 
             # Calculate the local integration error.
             local_error = local_error_estimation(

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -230,7 +230,7 @@ class RelaxedRigidContacts(ContactModel):
         )
 
         def _detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
-            x, y, z = jax.tree_map(jnp.squeeze, (x, y, z))
+            x, y, z = jax.tree.map(jnp.squeeze, (x, y, z))
 
             n̂ = self.terrain.normal(x=x, y=y).squeeze()
             h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])

jaxsim/utils/jaxsim_dataclass.py

@@ -298,7 +298,7 @@ class JaxsimDataclass(abc.ABC):
         """
 
         # Make a copy calling tree_map.
-        obj = jax.tree_util.tree_map(lambda leaf: leaf, self)
+        obj = jax.tree.map(lambda leaf: leaf, self)
 
         # Make sure that the copied object and all the copied leaves have the same
         # mutability of the original object.

{jaxsim-0.4.3.dev77.dist-info → jaxsim-0.4.3.dev88.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev77
+Version: 0.4.3.dev88
 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,9 +60,9 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
-Requires-Dist: jax>=0.4.13
+Requires-Dist: jax>=0.4.26
 Requires-Dist: jaxopt>=0.8.0
-Requires-Dist: jaxlib>=0.4.13
+Requires-Dist: jaxlib>=0.4.26
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax-dataclasses>=1.4.0
 Requires-Dist: pptree

{jaxsim-0.4.3.dev77.dist-info → jaxsim-0.4.3.dev88.dist-info}/RECORD

@@ -1,25 +1,25 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=iooIC7hyVsbupWXsxWXV06oZTPyCALANUUOu7UZTm8Y,426
+jaxsim/_version.py,sha256=B5nbhwRMLkROhGufRX4ovaMmwptyxOtLB8YwoD_nlo8,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 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=Ek1xSKB_VWjfqsqRYlK236ountKmGTl1M04cTYqHgsE,22142
+jaxsim/api/contact.py,sha256=d6uf_ySM5bqh90lHvxwcni722ePrvm3HMKy9Lo1uqqY,22017
 jaxsim/api/data.py,sha256=QldUHniJqKrdNtAcXuRaS9UyeslJ0Rjvb17UA0Ca5Tw,29008
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
-jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
-jaxsim/api/link.py,sha256=qPRtc8qqMRjZxUCZYXJMygbB6huDXBfIT1b1b8Durkw,18631
-jaxsim/api/model.py,sha256=TLjgacgTXm-2YRGDA0Id9pe9nxIem28KoAls6Tdk9WM,66241
+jaxsim/api/kin_dyn_parameters.py,sha256=ElahFk_RCcLvjTidH2qDOsY-m1gN1hXitCv4SvfgGYY,29260
+jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
+jaxsim/api/model.py,sha256=WL31JA2jK5L79TJ05ZouIYGe02rQvFItVDqizkzC1UE,66100
 jaxsim/api/ode.py,sha256=ZshGdHptftku0yoUwBiBdd1iOqntH0vVEOjRHfL7Fao,13518
 jaxsim/api/ode_data.py,sha256=7RSoBhfCJdP6P9InQbDwdBVpClPMMuetewI-6AWm-_0,20276
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=XIrJVJDO0ldaZ93WgoGNlFoRvazsRJTpO3DrK9kIXqM,20437
+jaxsim/integrators/common.py,sha256=78MBs89GxsL0wU2yAexjvBZt3HEtfZoGVIN9f0a8yTc,20305
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
-jaxsim/integrators/variable_step.py,sha256=5StkFh9oQba34zlkIoXG2fUN78gbxkHePWbrpQ-QZOI,21274
+jaxsim/integrators/variable_step.py,sha256=cJD98q5BaiSKvp_KY_1KN3PZpAUJR3L8YRmLX5WPJJo,21114
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
@@ -54,17 +54,17 @@ jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/rbda/contacts/common.py,sha256=VwAs742futAmLnDgbaOuLzNDBFiKDfYItdEZ4UcFgzE,2467
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=0rrk73MKMVPkEXf1MTEmCnXaPWzzEw6U9RHZtLNIu_E,13646
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=deTC0M2a_RER7iwVpxLCfuSlgBLqkTmHQdOJ4169IR4,13646
 jaxsim/rbda/contacts/rigid.py,sha256=zbSM0miwpgC1rp1d0RoQ1q8pYiKdIkHV8iZimeEPC94,15153
 jaxsim/rbda/contacts/soft.py,sha256=_wvb5iZDjGcVg6rNQelN4LZN7qSC2NIp0HdKvZmlGfk,15647
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=xUQg47yGxIOcTkLPbnO3sruEGBhoCd16j1evTGlmNjI,5010
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
-jaxsim/utils/jaxsim_dataclass.py,sha256=FSiUvdnq4Y1T9Jaa_mw4ZBQJe8H7deLr3Kupxtlh4iI,11322
+jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev77.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev77.dist-info/METADATA,sha256=SfRw58102asAL9OAPXqsvh9nr8FVvLamf9dCNL9XXPo,17276
-jaxsim-0.4.3.dev77.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev77.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev77.dist-info/RECORD,,
+jaxsim-0.4.3.dev88.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev88.dist-info/METADATA,sha256=Hm6Vv2DzMpviXdgHoSsUUuGWJAyYMKngkYXEVMHEh7U,17276
+jaxsim-0.4.3.dev88.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev88.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev88.dist-info/RECORD,,