Update README.md for Jax usage

README.md

@@ -16,14 +16,19 @@
 
 **adam** is based on Roy Featherstone's Rigid Body Dynamics Algorithms.
 
----
-
-<p align="center">
-  <b>⚠️ REPOSITORY UNDER DEVELOPMENT ⚠️</b>
-  <br>We cannot guarantee stable API
-</p>
-
----
+- [adam](#adam)
+  - [🐍 Dependencies](#-dependencies)
+  - [💾 Installation](#-installation)
+    - [🐍 Installation with pip](#-installation-with-pip)
+    - [📦 Installation with conda](#-installation-with-conda)
+      - [Installation from conda-forge package](#installation-from-conda-forge-package)
+    - [🔨 Installation from repo](#-installation-from-repo)
+  - [🚀 Usage](#-usage)
+    - [Jax interface](#jax-interface)
+    - [CasADi interface](#casadi-interface)
+    - [PyTorch interface](#pytorch-interface)
+  - [🦸‍♂️ Contributing](#️-contributing)
+  - [Todo](#todo)
 
 ## 🐍 Dependencies
 
@@ -158,6 +163,8 @@ import adam
 from adam.jax import KinDynComputations
 import icub_models
 import numpy as np
+import jax.numpy as jnp
+from jax import jit, vmap
 
 # if you want to icub-models https://github.com/robotology/icub-models to retrieve the urdf
 model_path = icub_models.get_model_file("iCubGazeboV2_5")
@@ -169,9 +176,8 @@ joints_name_list = [
     'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch',
     'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll'
 ]
-# Specify the root link
-root_link = 'root_link'
-kinDyn = KinDynComputations(model_path, joints_name_list, root_link)
+
+kinDyn = KinDynComputations(model_path, joints_name_list)
 # choose the representation, if you want to use the body fixed representation
 kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION)
 # or, if you want to use the mixed representation (that is the default)
@@ -180,6 +186,28 @@ w_H_b = np.eye(4)
 joints = np.ones(len(joints_name_list))
 M = kinDyn.mass_matrix(w_H_b, joints)
 print(M)
+w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b, joints)
+
+# IMPORTANT! The Jax Interface function execution can be slow! We suggest to jit them.
+# For example:
+
+def frame_forward_kinematics(w_H_b, joints):
+    # This is needed since str is not a valid JAX type
+    return kinDyn.forward_kinematics('frame_name', w_H_b, joints)
+
+jitted_frame_fk = jit(frame_forward_kinematics)
+w_H_f = jitted_frame_fk(w_H_b, joints)
+
+# In the same way, the functions can be also vmapped
+vmapped_frame_fk = vmap(frame_forward_kinematics, in_axes=(0, 0))
+# which can be also jitted
+jitted_vmapped_frame_fk = jit(vmapped_frame_fk)
+# and called on a batch of data
+joints_batch = jnp.tile(joints, (1024, 1))
+w_H_b_batch = jnp.tile(w_H_b, (1024, 1, 1))
+
+w_H_f_batch = jitted_vmapped_frame_fk(w_H_b_batch, joints_batch)
+
 ```
 
 ### CasADi interface
@@ -200,9 +228,8 @@ joints_name_list = [
     'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch',
     'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll'
 ]
-# Specify the root link
-root_link = 'root_link'
-kinDyn = KinDynComputations(model_path, joints_name_list, root_link)
+
+kinDyn = KinDynComputations(model_path, joints_name_list)
 # choose the representation you want to use the body fixed representation
 kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION)
 # or, if you want to use the mixed representation (that is the default)
@@ -245,9 +272,8 @@ joints_name_list = [
     'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch',
     'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll'
 ]
-# Specify the root link
-root_link = 'root_link'
-kinDyn = KinDynComputations(model_path, joints_name_list, root_link)
+
+kinDyn = KinDynComputations(model_path, joints_name_list)
 # choose the representation you want to use the body fixed representation
 kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION)
 # or, if you want to use the mixed representation (that is the default)
@@ -264,6 +290,9 @@ print(M)
 
 Open an issue with your feature request or if you spot a bug. Then, you can also proceed with a Pull-requests! :rocket:
 
+> [!WARNING]
+> REPOSITORY UNDER DEVELOPMENT! We cannot guarantee stable API
+
 ## Todo
 
 - [x] Center of Mass position