Repository: ami-iit/adam (language: python)

Issue:
Unify Jax and PyTorch backends 
It would be nice to unify the backends as done for CasADi and NumPy. 
I had to move some methods in the Jax computation class and reimplement them due to the immutable Jax types (no slice assign operator `[]`). The PyTorch implementation has a similar issue, since I had to cast some vectors in `torch,tensor`. 
Probably subclassing Jax and PyTorch and create a more abstract class in which some methods are redefined (for example `__setitem__` for handle the immutable types of jax, or a general `vector` type that casts array and list in `torch.tensor`) 


Discussion / hints:

Fix patch:
diff --git a/adam/casadi/spatial_math_casadi.py b/adam/casadi/spatial_math_casadi.py
deleted file mode 100644
index 6ce525fe..00000000
--- a/adam/casadi/spatial_math_casadi.py
+++ /dev/null
@@ -1,136 +0,0 @@
-import abc
-from typing import TypeVar
-
-import casadi as cs
-
-from adam.core.spatial_math import SpatialMathAbstract
-
-
-class SpatialMathCasadi(SpatialMathAbstract):
-    @classmethod
-    def R_from_axis_angle(cls, axis, q):
-        cq, sq = cs.cos(q), cs.sin(q)
-        return (
-            cq * (cls.eye(3) - cs.np.outer(axis, axis))
-            + sq * cls.skew(axis)
-            + cs.np.outer(axis, axis)
-        )
-
-    @classmethod
-    def Rx(cls, q):
-        R = cls.eye(3)
-        cq, sq = cs.cos(q), cs.sin(q)
-        R[1, 1] = cq
-        R[1, 2] = -sq
-        R[2, 1] = sq
-        R[2, 2] = cq
-        return R
-
-    @classmethod
-    def Ry(cls, q):
-        R = cls.eye(3)
-        cq, sq = cs.cos(q), cs.sin(q)
-        R[0, 0] = cq
-        R[0, 2] = sq
-        R[2, 0] = -sq
-        R[2, 2] = cq
-        return R
-
-    @classmethod
-    def Rz(cls, q):
-        R = cls.eye(3)
-        cq, sq = cs.cos(q), cs.sin(q)
-        R[0, 0] = cq
-        R[0, 1] = -sq
-        R[1, 0] = sq
-        R[1, 1] = cq
-        return R
-
-    @classmethod
-    def H_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.eye(4)
-        R = cls.R_from_RPY(rpy) @ cls.R_from_axis_angle(axis, q)
-        T[:3, :3] = R
-        T[:3, 3] = xyz
-        return T
-
-    @classmethod
-    def H_from_Pos_RPY(cls, xyz, rpy):
-        T = cls.eye(4)
-        T[:3, :3] = cls.R_from_RPY(rpy)
-        T[:3, 3] = xyz
-        return T
-
-    @classmethod
-    def R_from_RPY(cls, rpy):
-        return cls.Rz(rpy[2]) @ cls.Ry(rpy[1]) @ cls.Rx(rpy[0])
-
-    @classmethod
-    def X_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.H_revolute_joint(xyz, rpy, axis, q)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def X_fixed_joint(cls, xyz, rpy):
-        T = cls.H_from_Pos_RPY(xyz, rpy)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def spatial_transform(cls, R, p):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = R
-        X[3:, 3:] = R
-        X[:3, 3:] = cls.skew(p) @ R
-        return X
-
-    @classmethod
-    def spatial_inertia(cls, I, mass, c, rpy):
-        # Returns the 6x6 inertia matrix expressed at the origin of the link (with rotation)"""
-        IO = cls.zeros(6, 6)
-        Sc = cls.skew(c)
-        R = cls.R_from_RPY(rpy)
-        inertia_matrix = cs.np.array(
-            [[I.ixx, I.ixy, I.ixz], [I.ixy, I.iyy, I.iyz], [I.ixz, I.iyz, I.izz]]
-        )
-
-        IO[3:, 3:] = R @ inertia_matrix @ R.T + mass * Sc @ Sc.T
-        IO[3:, :3] = mass * Sc
-        IO[:3, 3:] = mass * Sc.T
-        IO[:3, :3] = cs.np.eye(3) * mass
-        return IO
-
-    @classmethod
-    def spatial_skew(cls, v):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = cls.skew(v[3:])
-        X[:3, 3:] = cls.skew(v[:3])
-        X[3:, 3:] = cls.skew(v[3:])
-        return X
-
-    @classmethod
-    def spatial_skew_star(cls, v):
-        return -cls.spatial_skew(v).T
-
-    @staticmethod
-    def zeros(*x):
-        return cs.SX.zeros(*x)
-
-    @staticmethod
-    def vertcat(*x):
-        return cs.vertcat(*x)
-
-    @staticmethod
-    def eye(x):
-        return cs.SX.eye(x)
-
-    @staticmethod
-    def skew(x):
-        return cs.skew(x)
-
-    @staticmethod
-    def array(*x):
-        return cs.SX.sym("vec", *x)
diff --git a/adam/core/spatial_math.py b/adam/core/spatial_math.py
deleted file mode 100644
index e0d2f8c9..00000000
--- a/adam/core/spatial_math.py
+++ /dev/null
@@ -1,92 +0,0 @@
-import abc
-
-import numpy as np
-import numpy.typing as npt
-
-
-class SpatialMathAbstract(abc.ABC):
-    @abc.abstractmethod
-    def R_from_axis_angle(self, axis: npt.NDArray, q: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def Rx(self, q: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def Ry(self, q: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def Rz(self, q: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def H_revolute_joint(
-        self,
-        xyz: npt.NDArray,
-        rpy: npt.NDArray,
-        axis: npt.NDArray,
-        q: npt.NDArray,
-    ):
-        pass
-
-    @abc.abstractmethod
-    def H_from_Pos_RPY(self, xyz: npt.NDArray, rpy: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def R_from_RPY(self, rpy: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def X_revolute_joint(
-        self,
-        xyz: npt.NDArray,
-        rpy: npt.NDArray,
-        axis: npt.NDArray,
-        q: npt.NDArray,
-    ):
-        pass
-
-    @abc.abstractmethod
-    def X_fixed_joint(self, xyz: npt.NDArray, rpy: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def spatial_transform(self, R: npt.NDArray, p: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def spatial_inertia(
-        self, I: npt.NDArray, mass: npt.NDArray, c: npt.NDArray, rpy: npt.NDArray
-    ):
-        pass
-
-    @abc.abstractmethod
-    def spatial_skew(self, v: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def spatial_skew_star(self, v: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def zeros(self, x: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def vertcat(self, x: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def eye(self, x: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def skew(self, x: npt.NDArray):
-        pass
-
-    @abc.abstractmethod
-    def array(self, x: npt.NDArray):
-        pass
diff --git a/adam/jax/computations.py b/adam/jax/computations.py
deleted file mode 100644
index c4d7c7e7..00000000
--- a/adam/jax/computations.py
+++ /dev/null
@@ -1,480 +0,0 @@
-# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
-# This software may be modified and distributed under the terms of the
-# GNU Lesser General Public License v2.1 or any later version.
-
-import jax.numpy as jnp
-import numpy as np
-from jax import grad, jit, vmap
-
-from adam.core.rbd_algorithms import RBDAlgorithms
-from adam.jax.spatial_math_jax import SpatialMathJax
-
-
-class KinDynComputations(RBDAlgorithms, SpatialMathJax):
-    """This is a small class that retrieves robot quantities using Jax
-    in mixed representation, for Floating Base systems - as humanoid robots.
-    """
-
-    def __init__(
-        self,
-        urdfstring: str,
-        joints_name_list: list,
-        root_link: str = "root_link",
-        gravity: np.array = jnp.array([0, 0, -9.80665, 0, 0, 0]),
-    ) -> None:
-        """
-        Args:
-            urdfstring (str): path of the urdf
-            joints_name_list (list): list of the actuated joints
-            root_link (str, optional): the first link. Defaults to 'root_link'.
-        """
-        super().__init__(
-            urdfstring=urdfstring,
-            joints_name_list=joints_name_list,
-            root_link=root_link,
-            gravity=gravity,
-        )
-
-    def mass_matrix(self, base_transform: jnp.array, joint_positions: jnp.array):
-        """Returns the Mass Matrix functions computed the CRBA
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            M (jax): Mass Matrix
-        """
-        [M, _] = self.crba(base_transform, joint_positions)
-        return M
-
-    def centroidal_momentum_matrix(
-        self, base_transform: jnp.array, joint_positions: jnp.array
-    ):
-        """Returns the Centroidal Momentum Matrix functions computed the CRBA
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            Jcc (jnp.ndarray): Centroidal Momentum matrix
-        """
-        [_, Jcm] = self.crba(base_transform, joint_positions)
-        return Jcm
-
-    def relative_jacobian(self, frame: str, joint_positions: jnp.array):
-        """Returns the Jacobian between the root link and a specified frame frames
-
-        Args:
-            frame (str): The tip of the chain
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            J (jnp.ndarray): The Jacobian between the root and the frame
-        """
-        return super().relative_jacobian(frame, joint_positions)
-
-    def CoM_position(self, base_transform: jnp.array, joint_positions: jnp.array):
-        """Returns the CoM positon
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            com (jnp.ndarray): The CoM position
-        """
-        return super().CoM_position(base_transform, joint_positions)
-
-    def crba(self, base_transform: jnp.array, joint_positions: jnp.array):
-        """This function computes the Composite Rigid body algorithm (Roy Featherstone) that computes the Mass Matrix.
-         The algorithm is complemented with Orin's modifications computing the Centroidal Momentum Matrix
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            M (jax): Mass Matrix
-            Jcm (jax): Centroidal Momentum Matrix
-        """
-        Ic = [None] * len(self.tree.links)
-        X_p = [None] * len(self.tree.links)
-        Phi = [None] * len(self.tree.links)
-        M = self.zeros(self.NDoF + 6, self.NDoF + 6)
-
-        for i in range(self.tree.N):
-            link_i = self.tree.links[i]
-            link_pi = self.tree.parents[i]
-            joint_i = self.tree.joints[i]
-            I = link_i.inertial.inertia
-            mass = link_i.inertial.mass
-            o = link_i.inertial.origin.xyz
-            rpy = link_i.inertial.origin.rpy
-            Ic[i] = self.spatial_inertia(I, mass, o, rpy)
-
-            if link_i.name == self.root_link:
-                # The first "real" link. The joint is universal.
-                X_p[i] = self.spatial_transform(np.eye(3), np.zeros(3))
-                Phi[i] = self.eye(6)
-            elif joint_i.type == "fixed":
-                X_J = self.X_fixed_joint(joint_i.origin.xyz, joint_i.origin.rpy)
-                X_p[i] = X_J
-                Phi[i] = self.zeros(6, 1)  # cs.vertcat(0, 0, 0, 0, 0, 0)
-            elif joint_i.type == "revolute" or joint_i.type == "continuous":
-                q_ = joint_positions[joint_i.idx] if joint_i.idx is not None else 0.0
-                X_J = self.X_revolute_joint(
-                    joint_i.origin.xyz,
-                    joint_i.origin.rpy,
-                    joint_i.axis,
-                    q_,
-                )
-                X_p[i] = X_J
-                Phi[i] = self.array(
-                    [
-                        0,
-                        0,
-                        0,
-                        joint_i.axis[0],
-                        joint_i.axis[1],
-                        joint_i.axis[2],
-                    ]
-                ).reshape(-1, 1)
-
-        for i in range(self.tree.N - 1, -1, -1):
-            link_i = self.tree.links[i]
-            link_pi = self.tree.parents[i]
-            joint_i = self.tree.joints[i]
-            if link_pi.name != self.tree.parents[0].name:
-                pi = self.tree.links.index(link_pi)
-                Ic[pi] = Ic[pi] + X_p[i].T @ Ic[i] @ X_p[i]
-            F = Ic[i] @ Phi[i]
-            if joint_i.idx is not None:
-                M = M.at[joint_i.idx + 6, joint_i.idx + 6].set((Phi[i].T @ F)[0, 0])
-            if link_i.name == self.root_link:
-                M = M.at[:6, :6].set(Phi[i].T @ F)
-            j = i
-            link_j = self.tree.links[j]
-            link_pj = self.tree.parents[j]
-            joint_j = self.tree.joints[j]
-            while self.tree.parents[j].name != self.tree.parents[0].name:
-                F = X_p[j].T @ F
-                j = self.tree.links.index(self.tree.parents[j])
-                joint_j = self.tree.joints[j]
-                if joint_i.name == self.tree.joints[0].name and joint_j.idx is not None:
-                    M = M.at[:6, joint_j.idx + 6].set(F.T @ Phi[j])
-                    M = M.at[joint_j.idx + 6, :6].set(M[:6, joint_j.idx + 6].T)
-                elif (
-                    joint_j.name == self.tree.joints[0].name and joint_i.idx is not None
-                ):
-                    M = M.at[joint_i.idx + 6, :6].set((F.T @ Phi[j])[0])
-                    M = M.at[:6, joint_i.idx + 6].set(M[joint_i.idx + 6, :6].T)
-                elif joint_i.idx is not None and joint_j.idx is not None:
-                    M = M.at[joint_i.idx + 6, joint_j.idx + 6].set((F.T @ Phi[j])[0, 0])
-                    M = M.at[joint_j.idx + 6, joint_i.idx + 6].set(
-                        M[joint_i.idx + 6, joint_j.idx + 6].T,
-                    )
-
-        X_G = [None] * len(self.tree.links)
-        O_X_G = self.eye(6)
-        O_X_G = O_X_G.at[:3, 3:].set(M[:3, 3:6].T / M[0, 0])
-        Jcm = self.zeros(6, self.NDoF + 6)
-        for i in range(self.tree.N):
-            link_i = self.tree.links[i]
-            link_pi = self.tree.parents[i]
-            joint_i = self.tree.joints[i]
-            if link_pi.name != self.tree.parents[0].name:
-                pi = self.tree.links.index(link_pi)
-                pi_X_G = X_G[pi]
-            else:
-                pi_X_G = O_X_G
-            X_G[i] = X_p[i] @ pi_X_G
-            if link_i.name == self.tree.links[0].name:
-                Jcm = Jcm.at[:, :6].set(X_G[i].T @ Ic[i] @ Phi[i])
-            elif joint_i.idx is not None:
-                Jcm = Jcm.at[:, joint_i.idx + 6].set((X_G[i].T @ Ic[i] @ Phi[i])[:, 0])
-
-        # Until now the algorithm returns the joint_positionsuantities in Body Fixed representation
-        # Moving to mixed representation...
-        X_to_mixed = self.eye(self.NDoF + 6)
-        X_to_mixed = X_to_mixed.at[:3, :3].set(base_transform[:3, :3].T)
-        X_to_mixed = X_to_mixed.at[3:6, 3:6].set(base_transform[:3, :3].T)
-        M = X_to_mixed.T @ M @ X_to_mixed
-        Jcc = X_to_mixed[:6, :6].T @ Jcm @ X_to_mixed
-        return M, Jcc
-
-    def forward_kinematics(
-        self, frame: str, base_transform: jnp.array, joint_positions: jnp.array
-    ):
-        """Computes the forward kinematics relative to the specified frame
-
-        Args:
-            frame (str): The frame to which the fk will be computed
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            T_fk (jnp.ndarray): The fk represented as Homogenous transformation matrix
-        """
-        return super().forward_kinematics(frame, base_transform, joint_positions)
-
-    def forward_kinematics_fun(self, frame):
-        fk_frame = lambda T, joint_positions: self.forward_kinematics(
-            frame, T, joint_positions
-        )
-        return fk_frame
-
-    def jacobian(self, frame, base_transform, joint_positions):
-        """Returns the Jacobian relative to the specified frame
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            s (jnp.ndarray): The joints position
-            frame (str): The frame to which the jacobian will be computed
-
-        Returns:
-            J_tot (jnp.ndarray): The Jacobian relative to the frame
-        """
-        chain = self.robot_desc.get_chain(self.root_link, frame)
-        T_fk = self.eye(4)
-        T_fk = T_fk @ base_transform
-        J = self.zeros(6, self.NDoF)
-        T_ee = self.forward_kinematics_fun(frame)
-        T_ee = T_ee(base_transform, joint_positions)
-        P_ee = T_ee[:3, 3]
-        for item in chain:
-            if item in self.robot_desc.joint_map:
-                joint = self.robot_desc.joint_map[item]
-                if joint.type == "fixed":
-                    xyz = joint.origin.xyz
-                    rpy = joint.origin.rpy
-                    joint_frame = self.H_from_Pos_RPY(xyz, rpy)
-                    T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
-                    T_joint = self.H_revolute_joint(
-                        joint.origin.xyz,
-                        joint.origin.rpy,
-                        joint.axis,
-                        q_,
-                    )
-                    T_fk = T_fk @ T_joint
-                    p_prev = P_ee - T_fk[:3, 3]
-                    z_prev = T_fk[:3, :3] @ jnp.array(joint.axis)
-                    if joint.idx is not None:
-                        J = J.at[:, joint.idx].set(
-                            jnp.vstack(
-                                [[self.skew(z_prev) @ p_prev], [z_prev]]
-                            ).reshape(-1),
-                        )
-
-        # Adding the floating base part of the Jacobian, in Mixed representation
-        J_tot = self.zeros(6, self.NDoF + 6)
-        J_tot = J_tot.at[:3, :3].set(self.eye(3))
-        J_tot = J_tot.at[:3, 3:6].set(-self.skew((P_ee - base_transform[:3, 3])))
-        J_tot = J_tot.at[:3, 6:].set(J[:3, :])
-        J_tot = J_tot.at[3:, 3:6].set(self.eye(3))
-        J_tot = J_tot.at[3:, 6:].set(J[3:, :])
-        return J_tot
-
-    def rnea(
-        self,
-        base_transform: jnp.array,
-        joint_positions: jnp.array,
-        base_velocity: jnp.array,
-        joint_velocities: jnp.array,
-        g: jnp.array,
-    ) -> jnp.array:
-        """Implementation of reduced Recursive Newton-Euler algorithm
-        (no acceleration and external forces). For now used to compute the bias force term
-
-        Args:
-            base_transform (jnp.array): The homogenous transform from base to world frame
-            joint_positions (jnp.array): The joints position
-            base_velocity (jnp.array): The base velocity in mixed representation
-            joint_velocities (jnp.array): The joints velocity
-            g (jnp.array): The 6D gravity acceleration
-
-        Returns:
-            tau (jnp.array): generalized force variables
-        """
-        # TODO: add accelerations
-        tau = self.zeros(self.NDoF + 6)
-        Ic = [None] * len(self.tree.links)
-        X_p = [None] * len(self.tree.links)
-        Phi = [None] * len(self.tree.links)
-        v = [None] * len(self.tree.links)
-        a = [None] * len(self.tree.links)
-        f = [None] * len(self.tree.links)
-
-        X_to_mixed = self.eye(6)
-        X_to_mixed = X_to_mixed.at[:3, :3].set(base_transform[:3, :3].T)
-        X_to_mixed = X_to_mixed.at[3:6, 3:6].set(base_transform[:3, :3].T)
-        acc_to_mixed = self.zeros(6, 1)
-        acc_to_mixed = acc_to_mixed.at[:3].set(
-            -base_transform[:3, :3].T
-            @ self.skew(base_velocity[3:])
-            @ base_velocity[:3].reshape(-1, 1),
-        )
-        acc_to_mixed = acc_to_mixed.at[3:].set(
-            -base_transform[:3, :3].T
-            @ self.skew(base_velocity[3:])
-            @ base_velocity[3:].reshape(-1, 1),
-        )
-        # set initial acceleration (rotated gravity + apparent acceleration)
-        a[0] = -X_to_mixed @ g.reshape(-1, 1) + acc_to_mixed
-
-        for i in range(self.tree.N):
-            link_i = self.tree.links[i]
-            link_pi = self.tree.parents[i]
-            joint_i = self.tree.joints[i]
-            I = link_i.inertial.inertia
-            mass = link_i.inertial.mass
-            o = link_i.inertial.origin.xyz
-            rpy = link_i.inertial.origin.rpy
-            Ic[i] = self.spatial_inertia(I, mass, o, rpy)
-
-            if link_i.name == self.root_link:
-                # The first "real" link. The joint is universal.
-                X_p[i] = self.spatial_transform(np.eye(3), np.zeros(3))
-                Phi[i] = self.eye(6)
-                v_J = Phi[i] @ X_to_mixed @ base_velocity.reshape(-1, 1)
-            elif joint_i.type == "fixed":
-                X_J = self.X_fixed_joint(joint_i.origin.xyz, joint_i.origin.rpy)
-                X_p[i] = X_J
-                Phi[i] = self.vertcat(0, 0, 0, 0, 0, 0)
-                v_J = self.zeros(6, 1)
-            elif joint_i.type == "revolute":
-                if joint_i.idx is not None:
-                    q_ = joint_positions[joint_i.idx]
-                    q_dot_ = joint_velocities[joint_i.idx]
-                else:
-                    q_ = 0.0
-                    q_dot_ = 0.0
-
-                X_J = self.X_revolute_joint(
-                    joint_i.origin.xyz,
-                    joint_i.origin.rpy,
-                    joint_i.axis,
-                    q_,
-                )
-                X_p[i] = X_J
-                Phi[i] = self.array(
-                    [0, 0, 0, joint_i.axis[0], joint_i.axis[1], joint_i.axis[2]]
-                ).reshape(-1, 1)
-                v_J = Phi[i] * q_dot_
-
-            if link_i.name == self.root_link:
-                v[i] = v_J
-                a[i] = X_p[i] @ a[0]
-            else:
-                pi = self.tree.links.index(link_pi)
-                v[i] = X_p[i] @ v[pi] + v_J
-                a[i] = X_p[i] @ a[pi] + self.spatial_skew(v[i]) @ v_J
-
-            f[i] = Ic[i] @ a[i] + self.spatial_skew_star(v[i]) @ Ic[i] @ v[i]
-
-        for i in range(self.tree.N - 1, -1, -1):
-            joint_i = self.tree.joints[i]
-            link_i = self.tree.links[i]
-            link_pi = self.tree.parents[i]
-            if joint_i.name == self.tree.joints[0].name:
-                tau = tau.at[:6].set((Phi[i].T @ f[i]).reshape(-1))
-            elif joint_i.idx is not None:
-                tau = tau.at[joint_i.idx + 6].set((Phi[i].T @ f[i])[0, 0])
-            if link_pi.name != self.tree.parents[0].name:
-                pi = self.tree.links.index(link_pi)
-                f[pi] = f[pi] + X_p[i].T @ f[i]
-        tau = tau.at[:6].set(X_to_mixed.T @ tau[:6])
-        return tau
-
-    def bias_force(
-        self,
-        base_transform: jnp.array,
-        joint_positions: jnp.array,
-        base_velocity: jnp.array,
-        s_dot: jnp.array,
-    ) -> jnp.array:
-        """Returns the bias force of the floating-base dynamics ejoint_positionsuation,
-        using a reduced RNEA (no acceleration and external forces)
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-            base_velocity (jnp.ndarray): The base velocity in mixed representation
-            s_dot (jnp.ndarray): The joints velocity
-
-        Returns:
-            h (jnp.ndarray): the bias force
-        """
-        h = self.rnea(base_transform, joint_positions, base_velocity, s_dot, self.g)
-        return h
-
-    def coriolis_term(
-        self,
-        base_transform: jnp.array,
-        joint_positions: jnp.array,
-        base_velocity: jnp.array,
-        s_dot: jnp.array,
-    ) -> jnp.array:
-        """Returns the coriolis term of the floating-base dynamics ejoint_positionsuation,
-        using a reduced RNEA (no acceleration and external forces)
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-            base_velocity (jnp.ndarray): The base velocity in mixed representation
-            s_dot (jnp.ndarray): The joints velocity
-
-        Returns:
-            C (jnp.ndarray): the Coriolis term
-        """
-        # set in the bias force computation the gravity term to zero
-        C = self.rnea(
-            base_transform,
-            joint_positions,
-            base_velocity.reshape(6, 1),
-            s_dot,
-            np.zeros(6),
-        )
-        return C
-
-    def gravity_term(
-        self, base_transform: jnp.array, joint_positions: jnp.array
-    ) -> jnp.array:
-        """Returns the gravity term of the floating-base dynamics ejoint_positionsuation,
-        using a reduced RNEA (no acceleration and external forces)
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            G (jnp.ndarray): the gravity term
-        """
-        G = self.rnea(
-            base_transform,
-            joint_positions,
-            np.zeros(6).reshape(6, 1),
-            np.zeros(self.NDoF),
-            self.g,
-        )
-        return G
-
-    def CoM_position(
-        self, base_transform: jnp.array, joint_positions: jnp.array
-    ) -> jnp.array:
-        """Returns the CoM positon
-
-        Args:
-            base_transform (jnp.ndarray): The homogenous transform from base to world frame
-            joint_positions (jnp.ndarray): The joints position
-
-        Returns:
-            com (jnp.ndarray): The CoM position
-        """
-        return super().CoM_position(base_transform, joint_positions)
diff --git a/adam/jax/spatial_math_jax.py b/adam/jax/spatial_math_jax.py
deleted file mode 100644
index 3fe9765d..00000000
--- a/adam/jax/spatial_math_jax.py
+++ /dev/null
@@ -1,149 +0,0 @@
-# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
-# This software may be modified and distributed under the terms of the
-# GNU Lesser General Public License v2.1 or any later version.
-
-import casadi as cs
-import jax.numpy as jnp
-import numpy as np
-
-from adam.core.spatial_math import SpatialMathAbstract
-
-
-class SpatialMathJax(SpatialMathAbstract):
-    @classmethod
-    def R_from_axis_angle(cls, axis, q):
-        cq, sq = jnp.cos(q), jnp.sin(q)
-        return (
-            cq * (jnp.eye(3) - jnp.outer(np.array(axis), np.array(axis)))
-            + sq * cls.skew(axis)
-            + jnp.outer(np.array(axis), np.array(axis))
-        )
-
-    @classmethod
-    def Rx(cls, q):
-        R = jnp.eye(3)
-        cq, sq = jnp.cos(q), jnp.sin(q)
-        R = R.at[1, 1].set(cq)
-        R = R.at[1, 2].set(-sq)
-        R = R.at[2, 1].set(sq)
-        R = R.at[2, 2].set(cq)
-        return R
-
-    @classmethod
-    def Ry(cls, q):
-        R = jnp.eye(3)
-        cq, sq = jnp.cos(q), jnp.sin(q)
-        R = R.at[0, 0].set(cq)
-        R = R.at[0, 2].set(sq)
-        R = R.at[2, 0].set(-sq)
-        R = R.at[2, 2].set(cq)
-        return R
-
-    @classmethod
-    def Rz(cls, q):
-        R = jnp.eye(3)
-        cq, sq = jnp.cos(q), jnp.sin(q)
-        R = R.at[0, 0].set(cq)
-        R = R.at[0, 1].set(-sq)
-        R = R.at[1, 0].set(sq)
-        R = R.at[1, 1].set(cq)
-        return R
-
-    @classmethod
-    def H_revolute_joint(cls, xyz, rpy, axis, q):
-        T = jnp.eye(4)
-        R = cls.R_from_RPY(rpy) @ cls.R_from_axis_angle(axis, q)
-        T = T.at[:3, :3].set(R)
-        T = T.at[:3, 3].set(xyz)
-        return T
-
-    @classmethod
-    def H_from_Pos_RPY(cls, xyz, rpy):
-        T = jnp.eye(4)
-        T = T.at[:3, :3].set(cls.R_from_RPY(rpy))
-        T = T.at[:3, 3].set(xyz)
-        return T
-
-    @classmethod
-    def R_from_RPY(cls, rpy):
-        return cls.Rz(rpy[2]) @ cls.Ry(rpy[1]) @ cls.Rx(rpy[0])
-
-    @classmethod
-    def X_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.H_revolute_joint(xyz, rpy, axis, q)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def X_fixed_joint(cls, xyz, rpy):
-        T = cls.H_from_Pos_RPY(xyz, rpy)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def spatial_transform(cls, R, p):
-        X = jnp.zeros([6, 6])
-        X = X.at[:3, :3].set(R)
-        X = X.at[3:, 3:].set(R)
-        X = X.at[:3, 3:].set(cls.skew(p) @ R)
-        return X
-
-    @classmethod
-    def spatial_inertia(cls, I, mass, c, rpy):
-        # Returns the 6x6 inertia matrix expressed at the origin of the link (with rotation)"""
-        IO = jnp.zeros([6, 6])
-        Sc = cls.skew(c)
-        R = cls.R_from_RPY(rpy)
-        inertia_matrix = np.array(
-            [[I.ixx, I.ixy, I.ixz], [I.ixy, I.iyy, I.iyz], [I.ixz, I.iyz, I.izz]]
-        )
-        IO = IO.at[:3, :3].set(jnp.eye(3) * mass)
-        IO = IO.at[3:, 3:].set(R @ inertia_matrix @ R.T + mass * Sc @ Sc.T)
-        IO = IO.at[3:, :3].set(mass * Sc)
-        IO = IO.at[:3, 3:].set(mass * Sc.T)
-        IO = IO.at[:3, :3].set(np.eye(3) * mass)
-        return IO
-
-    @classmethod
-    def spatial_skew(cls, v):
-        X = cls.zeros(6, 6)
-        X = X.at[:3, :3].set(cls.skew(v[3:]))
-        X = X.at[:3, 3:].set(cls.skew(v[:3]))
-        X = X.at[3:, 3:].set(cls.skew(v[3:]))
-        return X
-
-    @classmethod
-    def spatial_skew_star(cls, v):
-        return -cls.spatial_skew(v).T
-
-    @classmethod
-    def skew(cls, x):
-        return jnp.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
-
-    @staticmethod
-    def zeros(*x):
-        return jnp.zeros(x)
-
-    @staticmethod
-    def vertcat(*x):
-        v = jnp.vstack(x)
-        # This check is needed since vercat is used for two types of data structure in RBDAlgo class.
-        # CasADi handles the cases smootly, with NumPy I need to handle the two cases.
-        # It should be improved
-        if v.shape[1] > 1:
-            v = jnp.concatenate(x)
-        return v
-
-    @staticmethod
-    def eye(x):
-        return jnp.eye(x)
-
-    @staticmethod
-    def skew(x):
-        return -jnp.cross(jnp.array(x), jnp.eye(3), axisa=0, axisb=0)
-
-    @staticmethod
-    def array(*x):
-        return jnp.array(x)
diff --git a/adam/numpy/spatial_math_numpy.py b/adam/numpy/spatial_math_numpy.py
deleted file mode 100644
index 5f6e9d39..00000000
--- a/adam/numpy/spatial_math_numpy.py
+++ /dev/null
@@ -1,130 +0,0 @@
-from typing import TypeVar
-
-import numpy as np
-from scipy.spatial.transform import Rotation
-
-from adam.core.spatial_math import SpatialMathAbstract
-
-
-class SpatialMathNumpy(SpatialMathAbstract):
-    @classmethod
-    def R_from_axis_angle(cls, axis, q):
-        quat = cls.axis_angle_to_quat(axis, q)
-        return Rotation.from_quat(quat).as_matrix()
-
-    @classmethod
-    def axis_angle_to_quat(cls, axis, q):
-        s = np.sin(q / 2)
-        x = axis[0] * s
-        y = axis[1] * s
-        z = axis[2] * s
-        w = np.cos(q / 2)
-        return np.array([x, y, z, w])
-
-    @classmethod
-    def Rx(cls, q):
-        return Rotation.from_euler("x", q).as_matrix()
-
-    @classmethod
-    def Ry(cls, q):
-        return Rotation.from_euler("y", q).as_matrix()
-
-    @classmethod
-    def Rz(cls, q):
-        return Rotation.from_euler("z", q).as_matrix()
-
-    @classmethod
-    def H_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.eye(4)
-        R = cls.R_from_RPY(rpy) @ cls.R_from_axis_angle(axis, q)
-        T[:3, :3] = R
-        T[:3, 3] = xyz
-        return T
-
-    @classmethod
-    def H_from_Pos_RPY(cls, xyz, rpy):
-        T = cls.eye(4)
-        T[:3, :3] = cls.R_from_RPY(rpy)
-        T[:3, 3] = xyz
-        return T
-
-    @classmethod
-    def R_from_RPY(cls, rpy):
-        return Rotation.from_euler("xyz", rpy).as_matrix()
-
-    @classmethod
-    def X_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.H_revolute_joint(xyz, rpy, axis, q)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def X_fixed_joint(cls, xyz, rpy):
-        T = cls.H_from_Pos_RPY(xyz, rpy)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def spatial_transform(cls, R, p):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = R
-        X[3:, 3:] = R
-        X[:3, 3:] = cls.skew(p) @ R
-        return X
-
-    @classmethod
-    def spatial_inertia(cls, I, mass, c, rpy):
-        # Returns the 6x6 inertia matrix expressed at the origin of the link (with rotation)"""
-        IO = cls.zeros(6, 6)
-        Sc = cls.skew(c)
-        R = cls.R_from_RPY(rpy)
-        inertia_matrix = np.array(
-            [[I.ixx, I.ixy, I.ixz], [I.ixy, I.iyy, I.iyz], [I.ixz, I.iyz, I.izz]]
-        )
-
-        IO[3:, 3:] = R @ inertia_matrix @ R.T + mass * Sc @ Sc.T
-        IO[3:, :3] = mass * Sc
-        IO[:3, 3:] = mass * Sc.T
-        IO[:3, :3] = np.eye(3) * mass
-        return IO
-
-    @classmethod
-    def spatial_skew(cls, v):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = cls.skew(v[3:])
-        X[:3, 3:] = cls.skew(v[:3])
-        X[3:, 3:] = cls.skew(v[3:])
-        return X
-
-    @classmethod
-    def spatial_skew_star(cls, v):
-        return -cls.spatial_skew(v).T
-
-    @staticmethod
-    def zeros(*x):
-        return np.zeros(x)
-
-    @staticmethod
-    def vertcat(*x):
-        v = np.vstack(x)
-        # This check is needed since vercat is used for two types of data structure in RBDAlgo class.
-        # CasADi handles the cases smootly, with NumPy I need to handle the two cases.
-        # It should be improved
-        if v.shape[1] > 1:
-            v = np.concatenate(x)
-        return v
-
-    @staticmethod
-    def eye(x):
-        return np.eye(x)
-
-    @staticmethod
-    def skew(x):
-        # Retrieving the skew sym matrix using a cross product
-        return -np.cross(x, np.eye(3), axisa=0, axisb=0)
-
-    @staticmethod
-    def array(*x):
-        return np.empty(x)
diff --git a/adam/pytorch/spatial_math_pytorch.py b/adam/pytorch/spatial_math_pytorch.py
deleted file mode 100644
index d0d09f65..00000000
--- a/adam/pytorch/spatial_math_pytorch.py
+++ /dev/null
@@ -1,143 +0,0 @@
-import abc
-from typing import TypeVar
-
-import numpy as torch
-import numpy as np
-import torch
-
-from adam.core.spatial_math import SpatialMathAbstract
-
-
-class SpatialMathPytorch(SpatialMathAbstract):
-    @classmethod
-    def R_from_axis_angle(cls, axis, q):
-        q = torch.tensor(q)
-        cq, sq = torch.cos(q), torch.sin(q)
-        return (
-            cq * (cls.eye(3) - np.outer(axis, axis))
-            + sq * cls.skew(axis)
-            + np.outer(axis, axis)
-        )
-
-    @classmethod
-    def Rx(cls, q):
-        R = cls.eye(3)
-        q = torch.tensor(q)
-        cq, sq = torch.cos(q), torch.sin(q)
-        R[1, 1] = cq
-        R[1, 2] = -sq
-        R[2, 1] = sq
-        R[2, 2] = cq
-        return R
-
-    @classmethod
-    def Ry(cls, q):
-        R = cls.eye(3)
-        q = torch.tensor(q)
-        cq, sq = torch.cos(q), torch.sin(q)
-        R[0, 0] = cq
-        R[0, 2] = sq
-        R[2, 0] = -sq
-        R[2, 2] = cq
-        return R
-
-    @classmethod
-    def Rz(cls, q):
-        R = cls.eye(3)
-        q = torch.tensor(q)
-        cq, sq = torch.cos(q), torch.sin(q)
-        R[0, 0] = cq
-        R[0, 1] = -sq
-        R[1, 0] = sq
-        R[1, 1] = cq
-        return R
-
-    @classmethod
-    def H_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.eye(4)
-        R = cls.R_from_RPY(rpy).float() @ cls.R_from_axis_angle(axis, q).float()
-        T[:3, :3] = R
-        T[:3, 3] = torch.FloatTensor(xyz)
-        return T
-
-    @classmethod
-    def H_from_Pos_RPY(cls, xyz, rpy):
-        T = cls.eye(4)
-        T[:3, :3] = cls.R_from_RPY(rpy)
-        T[:3, 3] = torch.FloatTensor(xyz)
-        return T
-
-    @classmethod
-    def R_from_RPY(cls, rpy):
-        return cls.Rz(rpy[2]) @ cls.Ry(rpy[1]) @ cls.Rx(rpy[0])
-
-    @classmethod
-    def X_revolute_joint(cls, xyz, rpy, axis, q):
-        T = cls.H_revolute_joint(xyz, rpy, axis, q)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def X_fixed_joint(cls, xyz, rpy):
-        T = cls.H_from_Pos_RPY(xyz, rpy)
-        R = T[:3, :3].T
-        p = -T[:3, :3].T @ T[:3, 3]
-        return cls.spatial_transform(R, p)
-
-    @classmethod
-    def spatial_transform(cls, R, p):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = R
-        X[3:, 3:] = R
-        X[:3, 3:] = cls.skew(p) @ R
-        return X
-
-    @classmethod
-    def spatial_inertia(cls, I, mass, c, rpy):
-        # Returns the 6x6 inertia matrix expressed at the origin of the link (with rotation)"""
-        IO = cls.zeros(6, 6)
-        Sc = cls.skew(c)
-        R = cls.R_from_RPY(rpy)
-        inertia_matrix = torch.FloatTensor(
-            [[I.ixx, I.ixy, I.ixz], [I.ixy, I.iyy, I.iyz], [I.ixz, I.iyz, I.izz]]
-        )
-
-        IO[3:, 3:] = R @ inertia_matrix @ R.T + mass * Sc @ Sc.T
-        IO[3:, :3] = mass * Sc
-        IO[:3, 3:] = mass * Sc.T
-        IO[:3, :3] = torch.eye(3) * mass
-        return IO
-
-    @classmethod
-    def spatial_skew(cls, v):
-        X = cls.zeros(6, 6)
-        X[:3, :3] = cls.skew(v[3:])
-        X[:3, 3:] = cls.skew(v[:3])
-        X[3:, 3:] = cls.skew(v[3:])
-        return X
-
-    @classmethod
-    def spatial_skew_star(cls, v):
-        return -cls.spatial_skew(v).T
-
-    @staticmethod
-    def zeros(*x):
-        return torch.zeros(x).float()
-
-    @staticmethod
-    def vertcat(*x):
-        return torch.FloatTensor(x).reshape(-1, 1)
-
-    @staticmethod
-    def eye(x):
-        return torch.eye(x).float()
-
-    @staticmethod
-    def skew(x):
-        # Retrieving the skew sym matrix using a cross product
-        return torch.FloatTensor([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
-
-    @staticmethod
-    def array(*x):
-        return torch.empty(x)
diff --git a/examples/iCub_example_casadi.ipynb b/examples/iCub_example_casadi.ipynb
index 1da0633e..c2fe6739 100644
--- a/examples/iCub_example_casadi.ipynb
+++ b/examples/iCub_example_casadi.ipynb
@@ -2,17 +2,9 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 17,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from adam.casadi.computations import KinDynComputations\n",
     "from adam.geometry import utils\n",
@@ -21,16 +13,6 @@
     "import gym_ignition_models"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def SX2DM(sx):\n",
-    "    return cs.DM(sx)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -40,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -67,7 +49,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -89,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -98,41 +80,41 @@
      "text": [
       "Mass matrix:\n",
       " \n",
-      "[[33.0617, -2.37185e-07, -1.28839e-07, -5.36213e-08, -2.15628, -1.08778, -0.457875, -2.03524, 0.158567, -0.131484, 0.00919299, 0.0157066, 0.103774, 0.0337328, -0.0107114, 0.0055748, 0.0759185, -0.0300253, -0.212599, -0.0130793, 0.285848, -0.0277908, 0.00400305, 0.0142769, -0.135746, -0.095861, 0.35592, -0.0146894, 0.0292848], \n",
-      " [-2.37185e-07, 33.0617, -3.27854e-08, 2.15628, 3.70272e-08, 2.69101, -2.53357, 0.355518, -0.252342, 0.0298544, -0.281926, -0.0770974, 0.0126779, -0.0220408, -0.13744, 0.00403763, -0.0696663, -0.154873, 0.647578, 0.184583, 0.264584, -0.0227908, -0.0209135, 0.075277, -0.174444, -0.256781, -0.174708, 0.0104736, 0.0134755], \n",
-      " [-1.28839e-07, -3.27854e-08, 33.0617, 1.08778, -2.69101, 1.65941e-08, 0.34732, -1.13858, 0.138718, 0.25917, 0.0263738, -0.0200573, 0.0298214, 0.0090455, 0.376819, -0.0185187, 0.0345885, 0.353538, -0.0968578, -0.215565, 0.19638, -0.0167297, 0.0252119, -0.120089, -0.68269, -0.0131404, 0.147029, -0.0277721, 0.00706153], \n",
-      " [-5.36213e-08, 2.15628, 1.08778, 1.08188, -0.0661145, 0.144881, -0.714291, 0.171304, -0.205585, -0.0341696, -0.0176407, 0.0445861, -0.0103597, -0.0362648, -0.113403, 0.0228657, -0.0368076, 0.115892, -0.046504, -0.0853337, 0.0151912, -0.0006065, 0.011039, 0.0460579, -0.0227526, 0.0367208, 0.029983, -0.00805288, 0.000963516], \n",
-      " [-2.15628, 3.70272e-08, -2.69101, -0.0661145, 1.26892, -0.161349, 0.0871063, 0.730423, -0.00911797, -0.101164, -0.0128545, 0.0102908, 0.0226876, -0.0585365, 0.021172, 0.0324022, -0.0196076, -0.0316474, -0.0075284, 0.00588481, 0.0324758, -0.0111995, 0.00880397, 0.0195007, 0.209544, -0.0102159, 0.0140176, -0.00724137, 0.00298189], \n",
-      " [-1.08778, 2.69101, 1.65941e-08, 0.144881, -0.161349, 1.0107, -0.190237, -0.0454416, 0.233492, -0.00846543, -0.0973115, -0.0047234, 0.0260269, -0.009428, 0.00452672, 0.0133877, 0.0416197, 0.00314172, 0.145033, 0.0105944, -0.0661986, 0.0164676, 0.00560085, 0.0230043, -0.0412558, -0.0581033, -0.0544544, 0.00127344, -0.00986608], \n",
-      " [-0.457875, -2.53357, 0.34732, -0.714291, 0.0871063, -0.190237, 0.719563, -0.0390079, 0.199311, 0.0157703, 0.0156728, -0.0420306, 0.0140555, 0.025385, 0.11533, -0.0168484, 0.0325032, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [-2.03524, 0.355518, -1.13858, 0.171304, 0.730423, -0.0454416, -0.0390079, 0.683272, -0.0410906, -0.109794, -0.0129825, 0.0189947, 0.0234548, -0.0627546, 0.00122621, 0.0359632, -0.0234079, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.158567, -0.252342, 0.138718, -0.205585, -0.00911797, 0.233492, 0.199311, -0.0410906, 0.300322, 0.000914132, -0.0773358, -0.0288054, 0.0296679, 0.00549125, 0.0649337, 0.00203254, 0.0511466, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.131484, 0.0298544, 0.25917, -0.0341696, -0.101164, -0.00846543, 0.0157703, -0.109794, 0.000914132, 0.109197, 0.00157118, -0.0167064, -0.0384329, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.00919299, -0.281926, 0.0263738, -0.0176407, -0.0128545, -0.0973115, 0.0156728, -0.0129825, -0.0773358, 0.00157118, 0.0892657, 0.0121991, -0.017076, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.0157066, -0.0770974, -0.0200573, 0.0445861, 0.0102908, -0.0047234, -0.0420306, 0.0189947, -0.0288054, -0.0167064, 0.0121991, 0.0538243, -2.25892e-05, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.103774, 0.0126779, 0.0298214, -0.0103597, 0.0226876, 0.0260269, 0.0140555, 0.0234548, 0.0296679, -0.0384329, -0.017076, -2.25892e-05, 0.0464212, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.0337328, -0.0220408, 0.0090455, -0.0362648, -0.0585365, -0.009428, 0.025385, -0.0627546, 0.00549125, 0, 0, 0, 0, 0.0621118, 0.00412926, -0.0408066, 0.00679434, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.0107114, -0.13744, 0.376819, -0.113403, 0.021172, 0.00452672, 0.11533, 0.00122621, 0.0649337, 0, 0, 0, 0, 0.00412926, 0.141299, -0.0041829, 0.0318697, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.0055748, 0.00403763, -0.0185187, 0.0228657, 0.0324022, 0.0133877, -0.0168484, 0.0359632, 0.00203254, 0, 0, 0, 0, -0.0408066, -0.0041829, 0.0411333, -0.000755841, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [0.0759185, -0.0696663, 0.0345885, -0.0368076, -0.0196076, 0.0416197, 0.0325032, -0.0234079, 0.0511466, 0, 0, 0, 0, 0.00679434, 0.0318697, -0.000755841, 0.0464232, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.0300253, -0.154873, 0.353538, 0.115892, -0.0316474, 0.00314172, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0737865, -0.0140705, -0.0544491, -0.00424506, 0.00256794, 0.00580639, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.212599, 0.647578, -0.0968578, -0.046504, -0.0075284, 0.145033, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0140705, 0.147996, 0.0191458, -0.014742, 0.0102332, 0.00634051, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.0130793, 0.184583, -0.215565, -0.0853337, 0.00588481, 0.0105944, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0544491, 0.0191458, 0.0523891, 0.000822887, -0.000141661, -0.00736165, 0, 0, 0, 0, 0, 0], \n",
-      " [0.285848, 0.264584, 0.19638, 0.0151912, 0.0324758, -0.0661986, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.00424506, -0.014742, 0.000822887, 0.098885, -0.0217875, 0.000268157, 0, 0, 0, 0, 0, 0], \n",
-      " [-0.0277908, -0.0227908, -0.0167297, -0.0006065, -0.0111995, 0.0164676, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00256794, 0.0102332, -0.000141661, -0.0217875, 0.0138596, -1.4231e-05, 0, 0, 0, 0, 0, 0], \n",
-      " [0.00400305, -0.0209135, 0.0252119, 0.011039, 0.00880397, 0.00560085, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00580639, 0.00634051, -0.00736165, 0.000268157, -1.4231e-05, 0.0120266, 0, 0, 0, 0, 0, 0], \n",
-      " [0.0142769, 0.075277, -0.120089, 0.0460579, 0.0195007, 0.0230043, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0583148, 0.0133962, 0.0307806, 0.00654467, -0.00318644, 0.000128088], \n",
-      " [-0.135746, -0.174444, -0.68269, -0.0227526, 0.209544, -0.0412558, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0133962, 0.162484, 0.00961085, 0.0175254, -0.00928327, 0.00489682], \n",
-      " [-0.095861, -0.256781, -0.0131404, 0.0367208, -0.0102159, -0.0581033, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0307806, 0.00961085, 0.0486085, 0.00430511, -0.000854758, 0.00353587], \n",
-      " [0.35592, -0.174708, 0.147029, 0.029983, 0.0140176, -0.0544544, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00654467, 0.0175254, 0.00430511, 0.092591, -0.0182534, 0.00477585], \n",
-      " [-0.0146894, 0.0104736, -0.0277721, -0.00805288, -0.00724137, 0.00127344, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.00318644, -0.00928327, -0.000854758, -0.0182534, 0.0131613, -0.00042462], \n",
-      " [0.0292848, 0.0134755, 0.00706153, 0.000963516, 0.00298189, -0.00986608, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.000128088, 0.00489682, 0.00353587, 0.00477585, -0.00042462, 0.0120268]]\n"
+      "[[33.0617, 7.10849e-07, -3.95807e-07, -3.4761e-09, 0.0279637, -4.35541, 0.721991, -1.91601, 0.261959, -0.150573, -0.208839, 0.0567077, 0.0413244, 0.124995, -0.0831514, 0.0212593, 0.0532257, -0.344167, 0.711491, 0.217424, -0.142401, -0.0115969, 0.00375165, -0.17171, -0.0442504, 0.0316145, -0.417811, 0.0218527, -0.023327], \n",
+      " [7.10849e-07, 33.0617, -1.01395e-06, -0.0279638, -2.39614e-08, -0.581327, -0.383976, -1.50079, -0.0111384, -0.147059, -0.257414, 0.0661014, -0.0218793, 0.0924926, 0.234763, 0.0654011, 0.0103399, 0.0821493, -0.184482, -0.0519095, -0.328864, -0.0324675, -0.0119662, -0.123292, 0.801429, 0.236761, 0.000127647, -0.00741135, 0.0113491], \n",
+      " [-3.95807e-07, -1.01395e-06, 33.0617, 4.35541, 0.581327, 2.74375e-08, 1.39077, 0.356906, 0.253608, -0.317602, 0.194934, -0.0135158, 0.0981467, 0.0602712, 0.00354407, 0.0168676, -0.0941991, -0.701048, -0.412692, 0.213499, 0.0524024, 0.00287697, 0.030525, -0.411766, 0.823388, 0.314061, 0.0227567, 0.00548014, -0.0203987], \n",
+      " [-3.4761e-09, -0.0279638, 4.35541, 1.45615, -0.0586835, 0.23225, 0.436841, 0.31971, 0.158496, -0.0352486, 0.0756745, -0.045807, 0.0407701, 0.0240349, 0.0772524, 0.00587221, -0.00343431, -0.230636, -0.158699, 0.0564074, -0.0314623, 0.00453499, 0.00529636, -0.171667, 0.32854, 0.154392, 0.00757721, -0.00020314, -0.0123753], \n",
+      " [0.0279637, -2.39614e-08, 0.581327, -0.0586835, 0.860512, 0.179322, 0.243316, -0.235548, -0.20629, -0.155837, 0.041406, 0.0320002, 0.0616171, -0.0562398, 0.0317865, -0.0129744, 0.0276964, 0.0638904, -0.00759562, -0.0629583, 0.0129276, -0.000990284, 0.00596868, 0.0679094, -0.085742, -0.0364105, 0.0845898, -0.0129484, 0.00133846], \n",
+      " [-4.35541, -0.581327, 2.74375e-08, 0.23225, 0.179322, 1.44924, -0.125357, 0.586632, -0.195703, 0.0894609, 0.133815, -0.0348318, -0.00354522, 0.0303924, 0.0131741, 0.042736, 0.000980999, 0.113552, -0.257633, -0.0721661, -0.00584488, 0.00850546, 0.00173064, 0.0433781, 0.107385, 0.0116993, 0.14035, -0.0173388, 0.0149588], \n",
+      " [0.721991, -0.383976, 1.39077, 0.436841, 0.243316, -0.125357, 0.481568, 0.00188966, 0.191124, -0.0937845, 0.0179689, -0.021224, 0.0476879, -0.000208264, 0.0685327, -0.0147627, 0.000910508, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [-1.91601, -1.50079, 0.356906, 0.31971, -0.235548, 0.586632, 0.00188966, 0.774697, -0.0132183, 0.116535, 0.136426, -0.0629958, -0.00773232, 0.0534725, 0.0344464, 0.0414383, -0.0140909, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.261959, -0.0111384, 0.253608, 0.158496, -0.20629, -0.195703, 0.191124, -0.0132183, 0.330643, 0.0192834, -0.0573071, -0.0245243, -0.0113611, 0.0300205, 0.0115272, -0.0142202, -0.0206895, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.150573, -0.147059, -0.317602, -0.0352486, -0.155837, 0.0894609, -0.0937845, 0.116535, 0.0192834, 0.143702, 0.00420546, -0.0234813, -0.0531716, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.208839, -0.257414, 0.194934, 0.0756745, 0.041406, 0.133815, 0.0179689, 0.136426, -0.0573071, 0.00420546, 0.13057, -0.020941, 0.0248861, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.0567077, 0.0661014, -0.0135158, -0.045807, 0.0320002, -0.0348318, -0.021224, -0.0629958, -0.0245243, -0.0234813, -0.020941, 0.0491191, -0.00057869, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.0413244, -0.0218793, 0.0981467, 0.0407701, 0.0616171, -0.00354522, 0.0476879, -0.00773232, -0.0113611, -0.0531716, 0.0248861, -0.00057869, 0.0464212, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.124995, 0.0924926, 0.0602712, 0.0240349, -0.0562398, 0.0303924, -0.000208264, 0.0534725, 0.0300205, 0, 0, 0, 0, 0.0809855, 0.00228776, 0.0476742, -0.0143633, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.0831514, 0.234763, 0.00354407, 0.0772524, 0.0317865, 0.0131741, 0.0685327, 0.0344464, 0.0115272, 0, 0, 0, 0, 0.00228776, 0.0823856, 0.00552503, 0.0217545, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.0212593, 0.0654011, 0.0168676, 0.00587221, -0.0129744, 0.042736, -0.0147627, 0.0414383, -0.0142202, 0, 0, 0, 0, 0.0476742, 0.00552503, 0.0519112, 0.000236081, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [0.0532257, 0.0103399, -0.0941991, -0.00343431, 0.0276964, 0.000980999, 0.000910508, -0.0140909, -0.0206895, 0, 0, 0, 0, -0.0143633, 0.0217545, 0.000236081, 0.0464232, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.344167, 0.0821493, -0.701048, -0.230636, 0.0638904, 0.113552, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.161674, 0.0120213, -0.0425434, 0.00773452, -0.00202706, -0.0022285, 0, 0, 0, 0, 0, 0], \n",
+      " [0.711491, -0.184482, -0.412692, -0.158699, -0.00759562, -0.257633, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0120213, 0.194871, 0.0141275, 0.0206703, -0.0101046, -0.00171488, 0, 0, 0, 0, 0, 0], \n",
+      " [0.217424, -0.0519095, 0.213499, 0.0564074, -0.0629583, -0.0721661, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0425434, 0.0141275, 0.0535823, -0.00168219, -0.000300165, -0.00739459, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.142401, -0.328864, 0.0524024, -0.0314623, 0.0129276, -0.00584488, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00773452, 0.0206703, -0.00168219, 0.0686954, -0.00644098, 0.00229998, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.0115969, -0.0324675, 0.00287697, 0.00453499, -0.000990284, 0.00850546, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.00202706, -0.0101046, -0.000300165, -0.00644098, 0.0133563, 0.000436133, 0, 0, 0, 0, 0, 0], \n",
+      " [0.00375165, -0.0119662, 0.030525, 0.00529636, 0.00596868, 0.00173064, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0022285, -0.00171488, -0.00739459, 0.00229998, 0.000436133, 0.0120266, 0, 0, 0, 0, 0, 0], \n",
+      " [-0.17171, -0.123292, -0.411766, -0.171667, 0.0679094, 0.0433781, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.105788, -0.113972, -0.0824847, 0.0244279, -0.00396327, 0.0146364], \n",
+      " [-0.0442504, 0.801429, 0.823388, 0.32854, -0.085742, 0.107385, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.113972, 0.326092, 0.118636, 0.0155508, -0.00203121, 0.00249803], \n",
+      " [0.0316145, 0.236761, 0.314061, 0.154392, -0.0364105, 0.0116993, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0824847, 0.118636, 0.080658, -0.00294786, 0.00045249, -0.0128205], \n",
+      " [-0.417811, 0.000127647, 0.0227567, 0.00757721, 0.0845898, 0.14035, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0244279, 0.0155508, -0.00294786, 0.0902326, -0.0167133, 0.00502271], \n",
+      " [0.0218527, -0.00741135, 0.00548014, -0.00020314, -0.0129484, -0.0173388, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.00396327, -0.00203121, 0.00045249, -0.0167133, 0.0124395, -0.000669399], \n",
+      " [-0.023327, 0.0113491, -0.0203987, -0.0123753, 0.00133846, 0.0149588, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0146364, 0.00249803, -0.0128205, 0.00502271, -0.000669399, 0.0120268]]\n"
      ]
     }
    ],
    "source": [
     "M = kinDyn.mass_matrix_fun()\n",
-    "print('Mass matrix:\\n', SX2DM(M(H_b, s)))"
+    "print('Mass matrix:\\n', cs.DM(M(H_b, s)))"
    ]
   },
   {
@@ -144,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
@@ -153,18 +135,18 @@
      "text": [
       "Centroidal Momentum Matrix:\n",
       " \n",
-      "[[33.0617, -2.37185e-07, -1.28839e-07, -5.36213e-08, -2.15628, -1.08778, -0.457875, -2.03524, 0.158567, -0.131484, 0.00919299, 0.0157066, 0.103774, 0.0337328, -0.0107114, 0.0055748, 0.0759185, -0.0300253, -0.212599, -0.0130793, 0.285848, -0.0277908, 0.00400305, 0.0142769, -0.135746, -0.095861, 0.35592, -0.0146894, 0.0292848], \n",
-      " [-2.37185e-07, 33.0617, -3.27854e-08, 2.15628, 3.70272e-08, 2.69101, -2.53357, 0.355518, -0.252342, 0.0298544, -0.281926, -0.0770974, 0.0126779, -0.0220408, -0.13744, 0.00403763, -0.0696663, -0.154873, 0.647578, 0.184583, 0.264584, -0.0227908, -0.0209135, 0.075277, -0.174444, -0.256781, -0.174708, 0.0104736, 0.0134755], \n",
-      " [-1.28839e-07, -3.27854e-08, 33.0617, 1.08778, -2.69101, 1.65941e-08, 0.34732, -1.13858, 0.138718, 0.25917, 0.0263738, -0.0200573, 0.0298214, 0.0090455, 0.376819, -0.0185187, 0.0345885, 0.353538, -0.0968578, -0.215565, 0.19638, -0.0167297, 0.0252119, -0.120089, -0.68269, -0.0131404, 0.147029, -0.0277721, 0.00706153], \n",
-      " [-8.75344e-08, -3.62954e-08, -9.8257e-08, 0.905463, 0.022424, -0.0306267, -0.560479, 0.185578, -0.193691, -0.0446437, -0.000121249, 0.0502743, -0.0121678, -0.035125, -0.116837, 0.0232117, -0.033402, 0.114361, -0.0855521, -0.0902797, -0.00852611, 0.00143035, 0.0115735, 0.0450995, 0.0110861, 0.0539004, 0.03654, -0.00782223, -0.000147687], \n",
-      " [9.1163e-08, 5.59167e-08, 9.71107e-08, 0.022424, 0.909257, -0.232294, 0.0855135, 0.505011, 0.0125146, -0.0886446, -0.0101083, 0.00968267, 0.0318831, -0.0556002, 0.0511442, 0.0312585, -0.0118409, -0.00482989, -0.0292777, -0.0125139, 0.0671029, -0.0143737, 0.0111171, 0.0106573, 0.145124, -0.0175375, 0.0491978, -0.0104599, 0.0054666], \n",
-      " [-6.95873e-08, -1.63427e-08, 3.16177e-08, -0.0306267, -0.232294, 0.755874, 0.000914772, -0.141341, 0.259248, -0.0152214, -0.0740621, 0.00206862, 0.0284093, -0.00652416, 0.0153611, 0.0132425, 0.049788, 0.0147596, 0.0853296, -0.00485984, -0.0783293, 0.0174083, 0.00743479, 0.0173469, -0.0315234, -0.0403568, -0.0285239, -6.23491e-05, -0.00999938]]\n"
+      "[[33.0617, 7.10849e-07, -3.95807e-07, -3.4761e-09, 0.0279637, -4.35541, 0.721991, -1.91601, 0.261959, -0.150573, -0.208839, 0.0567077, 0.0413244, 0.124995, -0.0831514, 0.0212593, 0.0532257, -0.344167, 0.711491, 0.217424, -0.142401, -0.0115969, 0.00375165, -0.17171, -0.0442504, 0.0316145, -0.417811, 0.0218527, -0.023327], \n",
+      " [7.10849e-07, 33.0617, -1.01395e-06, -0.0279638, -2.39614e-08, -0.581327, -0.383976, -1.50079, -0.0111384, -0.147059, -0.257414, 0.0661014, -0.0218793, 0.0924926, 0.234763, 0.0654011, 0.0103399, 0.0821493, -0.184482, -0.0519095, -0.328864, -0.0324675, -0.0119662, -0.123292, 0.801429, 0.236761, 0.000127647, -0.00741135, 0.0113491], \n",
+      " [-3.95807e-07, -1.01395e-06, 33.0617, 4.35541, 0.581327, 2.74375e-08, 1.39077, 0.356906, 0.253608, -0.317602, 0.194934, -0.0135158, 0.0981467, 0.0602712, 0.00354407, 0.0168676, -0.0941991, -0.701048, -0.412692, 0.213499, 0.0524024, 0.00287697, 0.030525, -0.411766, 0.823388, 0.314061, 0.0227567, 0.00548014, -0.0203987], \n",
+      " [4.57909e-08, 1.0432e-07, -1.19258e-07, 0.882362, -0.135265, 0.231758, 0.253302, 0.271423, 0.125077, 0.00646662, 0.049777, -0.0439706, 0.0278222, 0.0161733, 0.0769841, 0.00370546, 0.00898383, -0.138213, -0.104489, 0.0282381, -0.0386437, 0.00412853, 0.00126501, -0.117527, 0.220748, 0.113219, 0.00457945, -0.000931339, -0.00967849], \n",
+      " [-2.38925e-08, -3.06957e-08, -5.54265e-08, -0.135265, 0.850267, 0.183006, 0.218251, -0.240203, -0.210971, -0.150125, 0.0381551, 0.0321899, 0.0598564, -0.0574053, 0.0317945, -0.0132889, 0.0293077, 0.0765081, -0.000940987, -0.0668961, 0.0121267, -0.00103106, 0.00542878, 0.0752947, -0.100182, -0.0419594, 0.0845431, -0.0130633, 0.00171686], \n",
+      " [2.52457e-08, 7.51251e-08, -1.50952e-08, 0.231758, 0.183006, 0.865259, -0.0369962, 0.307837, -0.161389, 0.0670393, 0.101777, -0.0261991, 0.00151398, 0.0484851, 0.0063479, 0.0466865, 0.00817454, 0.0696573, -0.167148, -0.0444363, -0.0303867, 0.00640686, 0.00201447, 0.0185899, 0.115647, 0.020027, 0.0853111, -0.0145903, 0.0120853]]\n"
      ]
     }
    ],
    "source": [
     "Jcm = kinDyn.centroidal_momentum_matrix_fun()\n",
-    "print('Centroidal Momentum Matrix:\\n', SX2DM(Jcm(H_b, s)))"
+    "print('Centroidal Momentum Matrix:\\n', cs.DM(Jcm(H_b, s)))"
    ]
   },
   {
@@ -176,7 +158,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
@@ -184,13 +166,13 @@
      "output_type": "stream",
      "text": [
       "Center of Mass position:\n",
-      " [-1.48223, -0.697022, -2.24606]\n"
+      " [2.39693, 1.03192, -2.08476]\n"
      ]
     }
    ],
    "source": [
     "CoM = kinDyn.CoM_position_fun()\n",
-    "print('Center of Mass position:\\n', SX2DM(CoM(H_b, s)))"
+    "print('Center of Mass position:\\n', cs.DM(CoM(H_b, s)))"
    ]
   },
   {
@@ -202,7 +184,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [
     {
@@ -227,7 +209,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -236,18 +218,18 @@
      "text": [
       "Jacobian of the left sole:\n",
       " \n",
-      "[[1, 00, 00, 0, 0.120413, -0.239225, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0166173, -0.0972176, -0.00661298, 0.141715, -0.0118362, 0.00388333, 0, 0, 0, 0, 0, 0], \n",
-      " [00, 1, 00, -0.120413, 0, -0.0301418, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0884937, 0.0162926, 0.0912776, 0.172288, -0.0497057, -0.0405297, 0, 0, 0, 0, 0, 0], \n",
-      " [00, 00, 1, 0.239225, 0.0301418, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.11404, -0.000408246, -0.111772, 0.130486, -0.0393521, 0.0500251, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 1, 00, 00, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.984185, 0.0218101, -0.998883, -0.0456993, 0.0456993, 0.183528, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 00, 1, 00, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.177039, 0.154888, -0.0380516, 0.626738, -0.626738, 0.770715, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 00, 00, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00603031, 0.987691, 0.0280244, -0.777889, 0.777889, 0.610177, 0, 0, 0, 0, 0, 0]]\n"
+      "[[1, 00, 00, 0, -0.143034, -0.261277, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.047121, 0.105796, 0.11144, -0.0855826, 0.00820937, 0.00337652, 0, 0, 0, 0, 0, 0], \n",
+      " [00, 1, 00, 0.143034, 0, 0.086364, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0441707, -0.0866635, -0.0137068, -0.137824, -0.045099, -0.0241324, 0, 0, 0, 0, 0, 0], \n",
+      " [00, 00, 1, 0.261277, -0.086364, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.116899, -0.0762706, 0.0598792, 0.0461249, -0.0186884, 0.05972, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 1, 00, 00, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.893528, -0.44666, 0.0163383, -0.671468, 0.671468, -0.165835, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 00, 1, 00, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.15855, 0.218337, -0.967615, 0.176054, -0.176054, 0.911032, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 00, 00, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.420083, -0.867654, -0.251902, -0.719816, 0.719816, 0.377518, 0, 0, 0, 0, 0, 0]]\n"
      ]
     }
    ],
    "source": [
     "J = kinDyn.jacobian_fun('l_sole')\n",
-    "print('Jacobian of the left sole:\\n', SX2DM(J(H_b, s)))"
+    "print('Jacobian of the left sole:\\n', cs.DM(J(H_b, s)))"
    ]
   },
   {
@@ -259,7 +241,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
@@ -268,18 +250,18 @@
      "text": [
       "Jacobian between the root link and left sole:\n",
       " \n",
-      "[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.136235, -0.000539507, 0.13352, -0.0635376, 0.0217028, -0.0596057, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2.82067e-14, -0.098567, -0.0233421, 0.109759, -0.00308654, 0.0112989, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0505218, 0.00106173, -0.0499579, -0.225186, 0.0606531, 0.021903, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2.07045e-13, -0.891505, -0.0982857, 0.926694, -0.926694, -0.336515, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1.8436e-13, -0.97618, -0.160624, 0.160624, 0.0478585, 0, 0, 0, 0, 0, 0], \n",
-      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.453011, 0.193422, -0.339762, 0.339762, -0.940461, 0, 0, 0, 0, 0, 0]]\n"
+      "[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.132545, -0.0689121, 0.0794147, -0.0493119, -0.0330415, 0.0405126, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2.74428e-14, -0.112831, -0.0722472, 0.117699, -0.00803553, 0.0258965, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.0163892, 0.0839095, 0.0683072, 0.110285, 0.0359758, 0.0429925, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2.07045e-13, -0.772787, -0.634319, -0.738473, 0.738473, 0.667461, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1.5978e-13, 0.0329968, 0.26968, -0.26968, 0.162323, 0, 0, 0, 0, 0, 0], \n",
+      " [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -0.634665, 0.772367, -0.618005, 0.618005, -0.726737, 0, 0, 0, 0, 0, 0]]\n"
      ]
     }
    ],
    "source": [
     "J_r = kinDyn.relative_jacobian_fun('l_sole')\n",
-    "print('Jacobian between the root link and left sole:\\n', SX2DM(J_r(s)))"
+    "print('Jacobian between the root link and left sole:\\n', cs.DM(J_r(s)))"
    ]
   },
   {
@@ -291,7 +273,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
@@ -300,16 +282,16 @@
      "text": [
       "Forward kinematics to the left sole:\n",
       " \n",
-      "[[0.183528, 0.0602066, 0.981169, -1.59376], \n",
-      " [0.770715, -0.628368, -0.105604, -0.490699], \n",
-      " [0.610177, 0.775583, -0.161725, -2.06043], \n",
+      "[[-0.165835, 0.0523491, 0.984763, 2.50087], \n",
+      " [0.911032, -0.374147, 0.173308, 1.16146], \n",
+      " [0.377518, 0.925891, 0.0143549, -2.22864], \n",
       " [0, 0, 0, 1]]\n"
      ]
     }
    ],
    "source": [
     "H = kinDyn.forward_kinematics_fun('l_sole')\n",
-    "print('Forward kinematics to the left sole:\\n', SX2DM(H(H_b, s)))"
+    "print('Forward kinematics to the left sole:\\n', cs.DM(H(H_b, s)))"
    ]
   },
   {
@@ -321,7 +303,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [
     {
@@ -329,13 +311,13 @@
      "output_type": "stream",
      "text": [
       "Bias force term:\n",
-      " [-10.4557, -12.6831, 342.605, 10.3614, -30.9899, 0.102771, 4.31766, -12.9124, 2.02004, 2.77223, 0.186025, -0.517362, 0.241313, 0.230028, 3.78126, -0.223227, 0.561087, 3.66779, -1.15611, -2.14063, 1.67583, -0.137267, 0.248258, -1.75708, -7.54027, 0.209693, 0.834118, -0.208901, 0.056177]\n"
+      " [16.1331, -51.9368, 339.587, 44.921, 8.13049, -1.27276, 17.3412, 5.09921, 2.64393, -4.29648, 3.01366, -0.281022, 1.94727, 0.758476, -0.336674, -0.0216558, -1.09433, -6.77529, -3.39509, 2.19611, 1.24952, 0.0694053, 0.321039, -3.99726, 6.46638, 2.54358, 0.064233, 0.128533, -0.257659]\n"
      ]
     }
    ],
    "source": [
     "h = kinDyn.bias_force_fun()\n",
-    "print('Bias force term:\\n', SX2DM(h(H_b, s, v_b, s_dot)))"
+    "print('Bias force term:\\n', cs.DM(h(H_b, s, v_b, s_dot)))"
    ]
   },
   {
@@ -347,7 +329,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 28,
    "metadata": {},
    "outputs": [
     {
@@ -355,13 +337,13 @@
      "output_type": "stream",
      "text": [
       "Coriolis force term:\n",
-      " [-10.4557, -12.6831, 18.3805, -0.306069, -4.60008, 0.102771, 0.911614, -1.74665, 0.659676, 0.230643, -0.0726137, -0.320667, -0.0511352, 0.141322, 0.0859238, -0.0416199, 0.22189, 0.200766, -0.206258, -0.0266594, -0.250005, 0.0267953, 0.00101389, -0.579415, -0.845361, 0.338557, -0.60774, 0.0634507, -0.013073]\n"
+      " [16.1332, -51.9368, 15.3625, 2.20902, 2.42962, -1.27276, 3.7024, 1.59915, 0.156886, -1.18186, 1.10201, -0.148477, 0.984778, 0.167417, -0.371429, -0.187071, -0.170549, 0.0996423, 0.652039, 0.102408, 0.735627, 0.0411919, 0.0216915, 0.0407844, -1.6083, -0.536308, -0.158934, 0.0747912, -0.0576157]\n"
      ]
     }
    ],
    "source": [
     "C = kinDyn.coriolis_term_fun()\n",
-    "print('Coriolis force term:\\n', SX2DM(C(H_b, s, v_b, s_dot)))"
+    "print('Coriolis force term:\\n', cs.DM(C(H_b, s, v_b, s_dot)))"
    ]
   },
   {
@@ -373,7 +355,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 29,
    "metadata": {},
    "outputs": [
     {
@@ -381,13 +363,13 @@
      "output_type": "stream",
      "text": [
       "Gravity term:\n",
-      " [-1.26348e-06, -3.21515e-07, 324.224, 10.6675, -26.3898, 1.62732e-07, 3.40605, -11.1657, 1.36036, 2.54158, 0.258638, -0.196695, 0.292448, 0.0887061, 3.69534, -0.181607, 0.339197, 3.46703, -0.94985, -2.11397, 1.92583, -0.164063, 0.247244, -1.17767, -6.69491, -0.128864, 1.44186, -0.272351, 0.06925]\n"
+      " [-3.88154e-06, -9.94347e-06, 324.224, 42.7119, 5.70087, 2.6907e-07, 13.6388, 3.50005, 2.48705, -3.11461, 1.91165, -0.132544, 0.96249, 0.591059, 0.0347554, 0.165415, -0.923778, -6.87493, -4.04713, 2.09371, 0.513892, 0.0282134, 0.299348, -4.03804, 8.07468, 3.07988, 0.223167, 0.0537418, -0.200043]\n"
      ]
     }
    ],
    "source": [
     "G = kinDyn.gravity_term_fun()\n",
-    "print('Gravity term:\\n', SX2DM(G(H_b, s)))"
+    "print('Gravity term:\\n', cs.DM(G(H_b, s)))"
    ]
   }
  ],
@@ -410,7 +392,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.12"
+   "version": "3.8.10"
   },
   "metadata": {
    "interpreter": {
diff --git a/pyproject.toml b/pyproject.toml
index 4cf57d88..e7db078a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,13 +1,16 @@
 [build-system]
 requires = [
-    "wheel",
-    "setuptools>=45",
-    "setuptools_scm[toml]>=6.0",
+  "wheel",
+  "setuptools>=45",
+  "setuptools_scm[toml]>=6.0",
 ]
 
+[tool.setuptools_scm]
+local_scheme = "dirty-tag"
+
 [tool.black]
 line-length = 88
 
 [tool.isort]
-profile = "black"
 multi_line_output = 3
+profile = "black"
diff --git a/setup.cfg b/setup.cfg
index 95f9d07b..a03823c0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -22,6 +22,8 @@ keywords =
 
 [options]
 packages = find:
+package_dir =
+        = src
 python_requires = >=3.8
 install_requires =
         jax
@@ -34,6 +36,9 @@ install_requires =
         urdf_parser_py
         torch
 
+[options.packages.find]
+where = src
+
 [options.extras_require]
 test =
         pytest
diff --git a/adam/__init__.py b/src/adam/__init__.py
similarity index 100%
rename from adam/__init__.py
rename to src/adam/__init__.py
diff --git a/adam/casadi/__init__.py b/src/adam/casadi/__init__.py
similarity index 100%
rename from adam/casadi/__init__.py
rename to src/adam/casadi/__init__.py
diff --git a/src/adam/casadi/casadi_like.py b/src/adam/casadi/casadi_like.py
new file mode 100644
index 00000000..3666ae06
--- /dev/null
+++ b/src/adam/casadi/casadi_like.py
@@ -0,0 +1,187 @@
+# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
+# This software may be modified and distributed under the terms of the
+# GNU Lesser General Public License v2.1 or any later version.
+
+from dataclasses import dataclass
+from typing import Union
+
+import casadi as cs
+import numpy.typing as npt
+
+from adam.core.spatial_math import ArrayLike
+
+
+@dataclass
+class CasadiLike(ArrayLike):
+    array: Union[cs.SX, cs.DM]
+
+    def __matmul__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array @ other.array)
+        else:
+            return CasadiLike(self.array @ other)
+
+    def __rmatmul__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return CasadiLike(other.array @ self.array)
+        else:
+            return CasadiLike(other @ self.array)
+
+    def __mul__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array * other.array)
+        else:
+            return CasadiLike(self.array * other)
+
+    def __rmul__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array * other.array)
+        else:
+            return CasadiLike(self.array * other)
+
+    def __add__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides + operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array + other.array)
+        else:
+            return CasadiLike(self.array + other)
+
+    def __radd__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides + operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array + other.array)
+        else:
+            return CasadiLike(self.array + other)
+
+    def __sub__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides - operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array - other.array)
+        else:
+            return CasadiLike(self.array - other)
+
+    def __rsub__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides - operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array - other.array)
+        else:
+            return CasadiLike(self.array - other)
+
+    def __neg__(self) -> "CasadiLike":
+        """Overrides - operator"""
+        return CasadiLike(-self.array)
+
+    def __truediv__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """Overrides / operator"""
+        if type(self) is type(other):
+            return CasadiLike(self.array / other.array)
+        else:
+            return CasadiLike(self.array / other)
+
+    def __setitem__(self, idx, value: Union["CasadiLike", npt.ArrayLike]):
+        """Overrides set item operator"""
+        self.array[idx] = value.array if type(self) is type(value) else value
+
+    def __getitem__(self, idx) -> "CasadiLike":
+        """Overrides get item operator"""
+        return CasadiLike(self.array[idx])
+
+    @property
+    def T(self) -> "CasadiLike":
+        """
+        Returns:
+            CasadiLike: Transpose of the array
+        """
+        return CasadiLike(self.array.T)
+
+    @staticmethod
+    def zeros(*x: int) -> "CasadiLike":
+        """
+        Returns:
+            CasadiLike: Matrix of zeros of dim *x
+        """
+        return CasadiLike(cs.SX.zeros(*x))
+
+    @staticmethod
+    def vertcat(*x) -> "CasadiLike":
+        """
+        Returns:
+            CasadiLike:  vertical concatenation of elements
+        """
+        # here the logic is a bit convoluted: x is a tuple containing CasadiLike
+        # cs.vertcat accepts *args. A list of cs types is created extracting the value
+        # from the CasadiLike stored in the tuple x.
+        # Then the list is unpacked with the * operator.
+        y = [xi.array if isinstance(xi, CasadiLike) else xi for xi in x]
+        return CasadiLike(cs.vertcat(*y))
+
+    @staticmethod
+    def eye(x: int) -> "CasadiLike":
+        """
+        Args:
+            x (int): matrix dimension
+
+        Returns:
+            CasadiLike: Identity matrix
+        """
+        return CasadiLike(cs.SX.eye(x))
+
+    @staticmethod
+    def skew(x: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """
+        Args:
+            x (Union[CasadiLike, npt.ArrayLike]): 3D vector
+
+        Returns:
+            CasadiLike: the skew symmetric matrix from x
+        """
+        if isinstance(x, CasadiLike):
+            return CasadiLike(cs.skew(x.array))
+        else:
+            return CasadiLike(cs.skew(x))
+
+    @staticmethod
+    def array(*x) -> "CasadiLike":
+        """
+        Returns:
+            CasadiLike: Vector wrapping *x
+        """
+        return CasadiLike(cs.DM(*x))
+
+    @staticmethod
+    def sin(x: npt.ArrayLike) -> "CasadiLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            CasadiLike: the sin value of x
+        """
+        return CasadiLike(cs.sin(x))
+
+    @staticmethod
+    def cos(x: npt.ArrayLike) -> "CasadiLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            CasadiLike: the cos value of x
+        """
+        return CasadiLike(cs.cos(x))
+
+    @staticmethod
+    def outer(x: npt.ArrayLike, y: npt.ArrayLike) -> "CasadiLike":
+        """
+        Args:
+            x (npt.ArrayLike): vector
+            y (npt.ArrayLike): vector
+
+        Returns:
+            CasadiLike: outer product between x and y
+        """
+        return CasadiLike(cs.np.outer(x, y))
diff --git a/adam/casadi/computations.py b/src/adam/casadi/computations.py
similarity index 87%
rename from adam/casadi/computations.py
rename to src/adam/casadi/computations.py
index f3eb115a..51253df1 100644
--- a/adam/casadi/computations.py
+++ b/src/adam/casadi/computations.py
@@ -5,11 +5,11 @@
 import casadi as cs
 import numpy as np
 
-from adam.casadi.spatial_math_casadi import SpatialMathCasadi
+from adam.casadi.casadi_like import CasadiLike
 from adam.core.rbd_algorithms import RBDAlgorithms
 
 
-class KinDynComputations(RBDAlgorithms, SpatialMathCasadi):
+class KinDynComputations(RBDAlgorithms, CasadiLike):
     """This is a small class that retrieves robot quantities represented in a symbolic fashion using CasADi
     in mixed representation, for Floating Base systems - as humanoid robots.
     """
@@ -45,7 +45,7 @@ def mass_matrix_fun(self) -> cs.Function:
         T_b = cs.SX.sym("T_b", 4, 4)
         s = cs.SX.sym("s", self.NDoF)
         [M, _] = super().crba(T_b, s)
-        return cs.Function("M", [T_b, s], [M], self.f_opts)
+        return cs.Function("M", [T_b, s], [M.array], self.f_opts)
 
     def centroidal_momentum_matrix_fun(self) -> cs.Function:
         """Returns the Centroidal Momentum Matrix functions computed the CRBA
@@ -56,7 +56,7 @@ def centroidal_momentum_matrix_fun(self) -> cs.Function:
         T_b = cs.SX.sym("T_b", 4, 4)
         s = cs.SX.sym("s", self.NDoF)
         [_, Jcm] = super().crba(T_b, s)
-        return cs.Function("Jcm", [T_b, s], [Jcm], self.f_opts)
+        return cs.Function("Jcm", [T_b, s], [Jcm.array], self.f_opts)
 
     def forward_kinematics_fun(self, frame: str) -> cs.Function:
         """Computes the forward kinematics relative to the specified frame
@@ -70,7 +70,7 @@ def forward_kinematics_fun(self, frame: str) -> cs.Function:
         s = cs.SX.sym("s", self.NDoF)
         T_b = cs.SX.sym("T_b", 4, 4)
         T_fk = super().forward_kinematics(frame, T_b, s)
-        return cs.Function("T_fk", [T_b, s], [T_fk], self.f_opts)
+        return cs.Function("T_fk", [T_b, s], [T_fk.array], self.f_opts)
 
     def jacobian_fun(self, frame: str) -> cs.Function:
         """Returns the Jacobian relative to the specified frame
@@ -84,7 +84,7 @@ def jacobian_fun(self, frame: str) -> cs.Function:
         s = cs.SX.sym("s", self.NDoF)
         T_b = cs.SX.sym("T_b", 4, 4)
         J_tot = super().jacobian(frame, T_b, s)
-        return cs.Function("J_tot", [T_b, s], [J_tot], self.f_opts)
+        return cs.Function("J_tot", [T_b, s], [J_tot.array], self.f_opts)
 
     def relative_jacobian_fun(self, frame: str) -> cs.Function:
         """Returns the Jacobian between the root link and a specified frame frames
@@ -97,7 +97,7 @@ def relative_jacobian_fun(self, frame: str) -> cs.Function:
         """
         s = cs.SX.sym("s", self.NDoF)
         J = super().relative_jacobian(frame, s)
-        return cs.Function("J", [s], [J], self.f_opts)
+        return cs.Function("J", [s], [J.array], self.f_opts)
 
     def CoM_position_fun(self) -> cs.Function:
         """Returns the CoM positon
@@ -108,7 +108,7 @@ def CoM_position_fun(self) -> cs.Function:
         s = cs.SX.sym("s", self.NDoF)
         T_b = cs.SX.sym("T_b", 4, 4)
         com_pos = super().CoM_position(T_b, s)
-        return cs.Function("CoM_pos", [T_b, s], [com_pos], self.f_opts)
+        return cs.Function("CoM_pos", [T_b, s], [com_pos.array], self.f_opts)
 
     def bias_force_fun(self) -> cs.Function:
         """Returns the bias force of the floating-base dynamics equation,
@@ -122,7 +122,7 @@ def bias_force_fun(self) -> cs.Function:
         v_b = cs.SX.sym("v_b", 6)
         s_dot = cs.SX.sym("s_dot", self.NDoF)
         h = super().rnea(T_b, s, v_b, s_dot, self.g)
-        return cs.Function("h", [T_b, s, v_b, s_dot], [h], self.f_opts)
+        return cs.Function("h", [T_b, s, v_b, s_dot], [h.array], self.f_opts)
 
     def coriolis_term_fun(self) -> cs.Function:
         """Returns the coriolis term of the floating-base dynamics equation,
@@ -137,7 +137,7 @@ def coriolis_term_fun(self) -> cs.Function:
         q_dot = cs.SX.sym("q_dot", self.NDoF)
         # set in the bias force computation the gravity term to zero
         C = super().rnea(T_b, q, v_b, q_dot, np.zeros(6))
-        return cs.Function("C", [T_b, q, v_b, q_dot], [C], self.f_opts)
+        return cs.Function("C", [T_b, q, v_b, q_dot], [C.array], self.f_opts)
 
     def gravity_term_fun(self) -> cs.Function:
         """Returns the gravity term of the floating-base dynamics equation,
@@ -150,4 +150,4 @@ def gravity_term_fun(self) -> cs.Function:
         q = cs.SX.sym("q", self.NDoF)
         # set in the bias force computation the velocity to zero
         G = super().rnea(T_b, q, np.zeros(6), np.zeros(self.NDoF), self.g)
-        return cs.Function("G", [T_b, q], [G], self.f_opts)
+        return cs.Function("G", [T_b, q], [G.array], self.f_opts)
diff --git a/adam/core/__init__.py b/src/adam/core/__init__.py
similarity index 100%
rename from adam/core/__init__.py
rename to src/adam/core/__init__.py
diff --git a/adam/core/rbd_algorithms.py b/src/adam/core/rbd_algorithms.py
similarity index 80%
rename from adam/core/rbd_algorithms.py
rename to src/adam/core/rbd_algorithms.py
index 92f1d6cd..b2ceab36 100644
--- a/adam/core/rbd_algorithms.py
+++ b/src/adam/core/rbd_algorithms.py
@@ -1,19 +1,13 @@
 # Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
 # This software may be modified and distributed under the terms of the
 # GNU Lesser General Public License v2.1 or any later version.
+import numpy.typing as npt
 
-
-from typing import TypeVar
-
-import numpy as np
-
-from adam.core.spatial_math import SpatialMathAbstract
+from adam.core.spatial_math import SpatialMath
 from adam.core.urdf_tree import URDFTree
 
-T = TypeVar("T")
 
-
-class RBDAlgorithms(SpatialMathAbstract):
+class RBDAlgorithms(SpatialMath):
     """This is a small abstract class that implements Rigid body algorithms retrieving robot quantities represented
     in mixed representation, for Floating Base systems - as humanoid robots.
     """
@@ -23,7 +17,7 @@ def __init__(
         urdfstring: str,
         joints_name_list: list,
         root_link: str,
-        gravity: np.array,
+        gravity: npt.ArrayLike,
     ) -> None:
         """
         Args:
@@ -47,16 +41,18 @@ def __init__(
             self.tree,
         ) = urdf_tree.load_model()
 
-    def crba(self, base_transform: T, joint_positions: T) -> T:
+    def crba(
+        self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
         """This function computes the Composite Rigid body algorithm (Roy Featherstone) that computes the Mass Matrix.
          The algorithm is complemented with Orin's modifications computing the Centroidal Momentum Matrix
 
         Args:
-            base_transform (T): The homogenous transform from base to world frame
-            joint_positions (T): The joints position
+            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
+            joint_positions (npt.ArrayLike): The joints position
         Returns:
-            M (T): Mass Matrix
-            Jcm (T): Centroidal Momentum Matrix
+            M (npt.ArrayLike): Mass Matrix
+            Jcm (npt.ArrayLike): Centroidal Momentum Matrix
         """
         Ic = [None] * len(self.tree.links)
         X_p = [None] * len(self.tree.links)
@@ -67,12 +63,8 @@ def crba(self, base_transform: T, joint_positions: T) -> T:
             link_i = self.tree.links[i]
             link_pi = self.tree.parents[i]
             joint_i = self.tree.joints[i]
-            I = link_i.inertial.inertia
-            mass = link_i.inertial.mass
-            o = link_i.inertial.origin.xyz
-            rpy = link_i.inertial.origin.rpy
+            I, mass, o, rpy = self.extract_link_properties(link_i)
             Ic[i] = self.spatial_inertia(I, mass, o, rpy)
-
             if link_i.name == self.root_link:
                 # The first "real" link. The joint is universal.
                 X_p[i] = self.spatial_transform(self.eye(3), self.zeros(3, 1))
@@ -81,16 +73,13 @@ def crba(self, base_transform: T, joint_positions: T) -> T:
                 X_J = self.X_fixed_joint(joint_i.origin.xyz, joint_i.origin.rpy)
                 X_p[i] = X_J
                 Phi[i] = self.vertcat(0, 0, 0, 0, 0, 0)
-            elif joint_i.type == "revolute" or joint_i.type == "continuous":
-                if joint_i.idx is not None:
-                    q_ = joint_positions[joint_i.idx]
-                else:
-                    q_ = 0.0
+            elif joint_i.type in ["revolute", "continuous"]:
+                q = joint_positions[joint_i.idx] if joint_i.idx is not None else 0.0
                 X_J = self.X_revolute_joint(
                     joint_i.origin.xyz,
                     joint_i.origin.rpy,
                     joint_i.axis,
-                    q_,
+                    q,
                 )
                 X_p[i] = X_J
                 Phi[i] = self.vertcat(
@@ -153,9 +142,9 @@ def crba(self, base_transform: T, joint_positions: T) -> T:
             if link_i.name == self.tree.links[0].name:
                 Jcm[:, :6] = X_G[i].T @ Ic[i] @ Phi[i]
             elif joint_i.idx is not None:
-                Jcm[:, [joint_i.idx + 6]] = X_G[i].T @ Ic[i] @ Phi[i]
+                Jcm[:, joint_i.idx + 6] = X_G[i].T @ Ic[i] @ Phi[i]
 
-        # Until now the algorithm returns the joint_positionsuantities in Body Fixed representation
+        # Until now the algorithm returns the joint_position quantities in Body Fixed representation
         # Moving to mixed representation...
         X_to_mixed = self.eye(self.NDoF + 6)
         X_to_mixed[:3, :3] = base_transform[:3, :3].T
@@ -164,15 +153,24 @@ def crba(self, base_transform: T, joint_positions: T) -> T:
         Jcm = X_to_mixed[:6, :6].T @ Jcm @ X_to_mixed
         return M, Jcm
 
-    def forward_kinematics(self, frame, base_transform: T, joint_positions: T) -> T:
+    def extract_link_properties(self, link_i):
+        I = link_i.inertial.inertia
+        mass = link_i.inertial.mass
+        o = link_i.inertial.origin.xyz
+        rpy = link_i.inertial.origin.rpy
+        return I, mass, o, rpy
+
+    def forward_kinematics(
+        self, frame, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
         """Computes the forward kinematics relative to the specified frame
 
         Args:
             frame (str): The frame to which the fk will be computed
-            base_transform (T): The homogenous transform from base to world frame
-            joint_positions (T): The joints position
+            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
+            joint_positions (npt.ArrayLike): The joints position
         Returns:
-            T_fk (T): The fk represented as Homogenous transformation matrix
+            T_fk (npt.ArrayLike): The fk represented as Homogenous transformation matrix
         """
         chain = self.robot_desc.get_chain(self.root_link, frame)
 
@@ -186,12 +184,9 @@ def forward_kinematics(self, frame, base_transform: T, joint_positions: T) -> T:
                     rpy = joint.origin.rpy
                     joint_frame = self.H_from_Pos_RPY(xyz, rpy)
                     T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
+                if joint.type in ["revolute", "continuous"]:
                     # if the joint is actuated set the value
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
+                    q_ = joint_positions[joint.idx] if joint.idx is not None else 0.0
                     T_joint = self.H_revolute_joint(
                         joint.origin.xyz,
                         joint.origin.rpy,
@@ -201,16 +196,18 @@ def forward_kinematics(self, frame, base_transform: T, joint_positions: T) -> T:
                     T_fk = T_fk @ T_joint
         return T_fk
 
-    def jacobian(self, frame: str, base_transform: T, joint_positions: T) -> T:
+    def jacobian(
+        self, frame: str, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
         """Returns the Jacobian relative to the specified frame
 
         Args:
             frame (str): The frame to which the jacobian will be computed
-            base_transform (T): The homogenous transform from base to world frame
-            joint_positions (T): The joints position
+            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
+            joint_positions (npt.ArrayLike): The joints position
 
         Returns:
-            J_tot (T): The Jacobian relative to the frame
+            J_tot (npt.ArrayLike): The Jacobian relative to the frame
         """
         chain = self.robot_desc.get_chain(self.root_link, frame)
         T_fk = self.eye(4)
@@ -226,11 +223,8 @@ def jacobian(self, frame: str, base_transform: T, joint_positions: T) -> T:
                     rpy = joint.origin.rpy
                     joint_frame = self.H_from_Pos_RPY(xyz, rpy)
                     T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
+                if joint.type in ["revolute", "continuous"]:
+                    q_ = joint_positions[joint.idx] if joint.idx is not None else 0.0
                     T_joint = self.H_revolute_joint(
                         joint.origin.xyz,
                         joint.origin.rpy,
@@ -238,7 +232,7 @@ def jacobian(self, frame: str, base_transform: T, joint_positions: T) -> T:
                         q_,
                     )
                     T_fk = T_fk @ T_joint
-                    p_prev = P_ee - T_fk[:3, 3]
+                    p_prev = P_ee - T_fk[:3, 3].array
                     z_prev = T_fk[:3, :3] @ joint.axis
                     # J[:, joint.idx] = self.vertcat(
                     #     cs.jacobian(P_ee, joint_positions[joint.idx]), z_prev) # using casadi jacobian
@@ -256,18 +250,20 @@ def jacobian(self, frame: str, base_transform: T, joint_positions: T) -> T:
         J_tot[3:, 6:] = J[3:, :]
         return J_tot
 
-    def relative_jacobian(self, frame: str, joint_positions: T) -> T:
+    def relative_jacobian(
+        self, frame: str, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
         """Returns the Jacobian between the root link and a specified frame frames
 
         Args:
             frame (str): The tip of the chain
-            joint_positions (T): The joints position
+            joint_positions (npt.ArrayLike): The joints position
 
         Returns:
-            J (T): The Jacobian between the root and the frame
+            J (npt.ArrayLike): The Jacobian between the root and the frame
         """
         chain = self.robot_desc.get_chain(self.root_link, frame)
-        base_transform = self.eye(4)
+        base_transform = self.eye(4).array
         T_fk = self.eye(4)
         T_fk = T_fk @ base_transform
         J = self.zeros(6, self.NDoF)
@@ -281,16 +277,13 @@ def relative_jacobian(self, frame: str, joint_positions: T) -> T:
                     rpy = joint.origin.rpy
                     joint_frame = self.H_from_Pos_RPY(xyz, rpy)
                     T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
+                if joint.type in ["revolute", "continuous"]:
+                    q = joint_positions[joint.idx] if joint.idx is not None else 0.0
                     T_joint = self.H_revolute_joint(
                         joint.origin.xyz,
                         joint.origin.rpy,
                         joint.axis,
-                        q_,
+                        q,
                     )
                     T_fk = T_fk @ T_joint
                     p_prev = P_ee - T_fk[:3, 3]
@@ -303,7 +296,9 @@ def relative_jacobian(self, frame: str, joint_positions: T) -> T:
                         )
         return J
 
-    def CoM_position(self, base_transform: T, joint_positions: T) -> T:
+    def CoM_position(
+        self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
         """Returns the CoM positon
 
         Args:
@@ -313,7 +308,7 @@ def CoM_position(self, base_transform: T, joint_positions: T) -> T:
         Returns:
             com (T): The CoM position
         """
-        com_pos = self.zeros(3)
+        com_pos = self.zeros(3, 1)
         for item in self.robot_desc.link_map:
             link = self.robot_desc.link_map[item]
             if link.inertial is not None:
@@ -323,7 +318,7 @@ def CoM_position(self, base_transform: T, joint_positions: T) -> T:
                     link.inertial.origin.rpy,
                 )
                 # Adding the link transform
-                T_fk = T_fk @ T_link
+                T_fk = T_fk @ T_link.array
                 com_pos += T_fk[:3, 3] * link.inertial.mass
         com_pos /= self.get_total_mass()
         return com_pos
@@ -343,12 +338,12 @@ def get_total_mass(self):
 
     def rnea(
         self,
-        base_transform: T,
-        joint_positions: T,
-        base_velocity: T,
-        joint_velocities: T,
-        g: T,
-    ) -> T:
+        base_transform: npt.ArrayLike,
+        joint_positions: npt.ArrayLike,
+        base_velocity: npt.ArrayLike,
+        joint_velocities: npt.ArrayLike,
+        g: npt.ArrayLike,
+    ) -> npt.ArrayLike:
         """Implementation of reduced Recursive Newton-Euler algorithm
         (no acceleration and external forces). For now used to compute the bias force term
 
@@ -378,10 +373,10 @@ def rnea(
 
         acc_to_mixed = self.zeros(6, 1)
         acc_to_mixed[:3] = (
-            -base_transform[:3, :3].T @ self.skew(base_velocity[3:]) @ base_velocity[:3]
+            -X_to_mixed[:3, :3] @ self.skew(base_velocity[3:]) @ base_velocity[:3]
         )
         acc_to_mixed[3:] = (
-            -base_transform[:3, :3].T @ self.skew(base_velocity[3:]) @ base_velocity[3:]
+            -X_to_mixed[:3, :3] @ self.skew(base_velocity[3:]) @ base_velocity[3:]
         )
         # set initial acceleration (rotated gravity + apparent acceleration)
         # reshape g as a vertical vector
@@ -391,10 +386,7 @@ def rnea(
             link_i = self.tree.links[i]
             link_pi = self.tree.parents[i]
             joint_i = self.tree.joints[i]
-            I = link_i.inertial.inertia
-            mass = link_i.inertial.mass
-            o = link_i.inertial.origin.xyz
-            rpy = link_i.inertial.origin.rpy
+            I, mass, o, rpy = self.extract_link_properties(link_i)
             Ic[i] = self.spatial_inertia(I, mass, o, rpy)
 
             if link_i.name == self.root_link:
@@ -407,25 +399,22 @@ def rnea(
                 X_p[i] = X_J
                 Phi[i] = self.vertcat(0, 0, 0, 0, 0, 0)
                 v_J = self.zeros(6, 1)
-            elif joint_i.type == "revolute" or joint_i.type == "continuous":
-                if joint_i.idx is not None:
-                    q_ = joint_positions[joint_i.idx]
-                    joint_velocities_ = joint_velocities[joint_i.idx]
-                else:
-                    q_ = 0.0
-                    joint_velocities_ = 0.0
-
+            elif joint_i.type in ["revolute", "continuous"]:
+                q = joint_positions[joint_i.idx] if joint_i.idx is not None else 0.0
+                q_dot = (
+                    joint_velocities[joint_i.idx] if joint_i.idx is not None else 0.0
+                )
                 X_J = self.X_revolute_joint(
                     joint_i.origin.xyz,
                     joint_i.origin.rpy,
                     joint_i.axis,
-                    q_,
+                    q,
                 )
                 X_p[i] = X_J
                 Phi[i] = self.vertcat(
                     0, 0, 0, joint_i.axis[0], joint_i.axis[1], joint_i.axis[2]
                 )
-                v_J = Phi[i] * joint_velocities_
+                v_J = Phi[i] * q_dot
 
             if link_i.name == self.root_link:
                 v[i] = v_J
diff --git a/src/adam/core/spatial_math.py b/src/adam/core/spatial_math.py
new file mode 100644
index 00000000..15156900
--- /dev/null
+++ b/src/adam/core/spatial_math.py
@@ -0,0 +1,299 @@
+import abc
+
+import numpy as np
+import numpy.typing as npt
+
+
+class ArrayLike(abc.ABC):
+    """Abstract class for a generic Array wrapper. Every method should be implemented for every data type."""
+
+    @abc.abstractmethod
+    def zeros(x: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            x (npt.ArrayLike): matrix dimension
+
+        Returns:
+            npt.ArrayLike: zero matrix of dimension x
+        """
+        pass
+
+    @abc.abstractmethod
+    def vertcat(x: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            x (npt.ArrayLike): elements
+
+        Returns:
+            npt.ArrayLike: vertical concatenation of elements x
+        """
+        pass
+
+    @abc.abstractmethod
+    def eye(x: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            x (npt.ArrayLike): matrix dimension
+
+        Returns:
+            npt.ArrayLike: identity matrix of dimension x
+        """
+        pass
+
+    @abc.abstractmethod
+    def sin(x: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: sin value of x
+        """
+        pass
+
+    @abc.abstractmethod
+    def cos(x: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: cos value of angle x
+        """
+        pass
+
+
+class SpatialMath(ArrayLike):
+    """Class implementing the main geometric functions used for computing rigid-body algorithm
+
+    Args:
+        ArrayLike: abstract class describing a generic Array wrapper. It needs to be implemented for every data type
+
+    """
+
+    @classmethod
+    def R_from_axis_angle(cls, axis: npt.ArrayLike, q: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            axis (npt.ArrayLike): axis vector
+            q (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: rotation matrix from axis-angle representation
+        """
+        cq, sq = cls.cos(q), cls.sin(q)
+        return (
+            cq * (cls.eye(3) - cls.outer(axis, axis))
+            + sq * cls.skew(axis)
+            + cls.outer(axis, axis)
+        )
+
+    @classmethod
+    def Rx(cls, q: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            q (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: rotation matrix around x axis
+        """
+        R = cls.eye(3)
+        cq, sq = cls.cos(q), cls.sin(q)
+        R[1, 1] = cq
+        R[1, 2] = -sq
+        R[2, 1] = sq
+        R[2, 2] = cq
+        return R
+
+    @classmethod
+    def Ry(cls, q: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            q (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: rotation matrix around y axis
+        """
+        R = cls.eye(3)
+        cq, sq = cls.cos(q), cls.sin(q)
+        R[0, 0] = cq
+        R[0, 2] = sq
+        R[2, 0] = -sq
+        R[2, 2] = cq
+        return R
+
+    @classmethod
+    def Rz(cls, q: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            q (npt.ArrayLike): angle value
+
+        Returns:
+            npt.ArrayLike: rotation matrix around z axis
+        """
+        R = cls.eye(3)
+        cq, sq = cls.cos(q), cls.sin(q)
+        R[0, 0] = cq
+        R[0, 1] = -sq
+        R[1, 0] = sq
+        R[1, 1] = cq
+        return R
+
+    @classmethod
+    def H_revolute_joint(
+        cls,
+        xyz: npt.ArrayLike,
+        rpy: npt.ArrayLike,
+        axis: npt.ArrayLike,
+        q: npt.ArrayLike,
+    ) -> npt.ArrayLike:
+        """
+        Args:
+            xyz (npt.ArrayLike): joint origin in the urdf
+            rpy (npt.ArrayLike): joint orientation in the urdf
+            axis (npt.ArrayLike): joint axis in the urdf
+            q (npt.ArrayLike): joint angle value
+
+        Returns:
+            npt.ArrayLike: Homogeneous transform
+        """
+        T = cls.eye(4)
+        R = cls.R_from_RPY(rpy) @ cls.R_from_axis_angle(axis, q)
+        T[:3, :3] = R
+        T[:3, 3] = xyz
+        return T
+
+    @classmethod
+    def H_from_Pos_RPY(cls, xyz: npt.ArrayLike, rpy: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            xyz (npt.ArrayLike): translation vector
+            rpy (npt.ArrayLike): rotation as rpy angles
+
+        Returns:
+            npt.ArrayLike: Homegeneous transform
+        """
+        T = cls.eye(4)
+        T[:3, :3] = cls.R_from_RPY(rpy)
+        T[:3, 3] = xyz
+        return T
+
+    @classmethod
+    def R_from_RPY(cls, rpy: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+           rpy (npt.ArrayLike): rotation as rpy angles
+
+        Returns:
+            npt.ArrayLike: Rotation matrix
+        """
+        return cls.Rz(rpy[2]) @ cls.Ry(rpy[1]) @ cls.Rx(rpy[0])
+
+    @classmethod
+    def X_revolute_joint(
+        cls,
+        xyz: npt.ArrayLike,
+        rpy: npt.ArrayLike,
+        axis: npt.ArrayLike,
+        q: npt.ArrayLike,
+    ) -> npt.ArrayLike:
+        """
+        Args:
+            xyz (npt.ArrayLike): joint origin in the urdf
+            rpy (npt.ArrayLike): joint orientation in the urdf
+            axis (npt.ArrayLike): joint axis in the urdf
+            q (npt.ArrayLike): joint angle value
+
+        Returns:
+            npt.ArrayLike: Spatial transform of a revolute joint given its rotation angle
+        """
+        # TODO: give Featherstone reference
+        T = cls.H_revolute_joint(xyz, rpy, axis, q)
+        R = T[:3, :3].T
+        p = -T[:3, :3].T @ T[:3, 3]
+        return cls.spatial_transform(R, p)
+
+    @classmethod
+    def X_fixed_joint(cls, xyz: npt.ArrayLike, rpy: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            xyz (npt.ArrayLike): joint origin in the urdf
+            rpy (npt.ArrayLike): joint orientation in the urdf
+
+        Returns:
+            npt.ArrayLike: Spatial transform of a fixed joint
+        """
+        T = cls.H_from_Pos_RPY(xyz, rpy)
+        R = T[:3, :3].T
+        p = -T[:3, :3].T @ T[:3, 3]
+        return cls.spatial_transform(R, p)
+
+    @classmethod
+    def spatial_transform(cls, R: npt.ArrayLike, p: npt.ArrayLike) -> npt.ArrayLike:
+        """_summary_
+
+        Args:
+            R (npt.ArrayLike): Rotation matrix
+            p (npt.ArrayLike): translation vector
+
+        Returns:
+            npt.ArrayLike: spatial transform
+        """
+        X = cls.zeros(6, 6)
+        X[:3, :3] = R
+        X[3:, 3:] = R
+        X[:3, 3:] = cls.skew(p) @ R
+        return X
+
+    @classmethod
+    def spatial_inertia(
+        cls, I: npt.ArrayLike, mass: npt.ArrayLike, c: npt.ArrayLike, rpy: npt.ArrayLike
+    ) -> npt.ArrayLike:
+        """
+        Args:
+            I (npt.ArrayLike): inertia values from urdf
+            mass (npt.ArrayLike): mass value from urdf
+            c (npt.ArrayLike): origin of the link from urdf
+            rpy (npt.ArrayLike): orientation of the link from the urdf
+
+        Returns:
+            npt.ArrayLike: the 6x6 inertia matrix expressed at the origin of the link (with rotation)
+        """
+        IO = cls.zeros(6, 6)
+        Sc = cls.skew(c)
+        R = cls.R_from_RPY(rpy)
+        inertia_matrix = cls.array(
+            [[I.ixx, I.ixy, I.ixz], [I.ixy, I.iyy, I.iyz], [I.ixz, I.iyz, I.izz]]
+        )
+
+        IO[3:, 3:] = R @ inertia_matrix @ R.T + mass * Sc @ Sc.T
+        IO[3:, :3] = mass * Sc
+        IO[:3, 3:] = mass * Sc.T
+        IO[:3, :3] = cls.eye(3) * mass
+        return IO
+
+    @classmethod
+    def spatial_skew(cls, v: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            v (npt.ArrayLike): 6D vector
+
+        Returns:
+            npt.ArrayLike: spatial skew matrix
+        """
+        X = cls.zeros(6, 6)
+        X[:3, :3] = cls.skew(v[3:])
+        X[:3, 3:] = cls.skew(v[:3])
+        X[3:, 3:] = cls.skew(v[3:])
+        return X
+
+    @classmethod
+    def spatial_skew_star(cls, v: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            v (npt.ArrayLike): 6D vector
+
+        Returns:
+            npt.ArrayLike: negative spatial skew matrix traspose
+        """
+        return -cls.spatial_skew(v).T
diff --git a/adam/core/urdf_tree.py b/src/adam/core/urdf_tree.py
similarity index 100%
rename from adam/core/urdf_tree.py
rename to src/adam/core/urdf_tree.py
diff --git a/adam/geometry/__init__.py b/src/adam/geometry/__init__.py
similarity index 100%
rename from adam/geometry/__init__.py
rename to src/adam/geometry/__init__.py
diff --git a/adam/geometry/utils.py b/src/adam/geometry/utils.py
similarity index 100%
rename from adam/geometry/utils.py
rename to src/adam/geometry/utils.py
diff --git a/adam/jax/__init__.py b/src/adam/jax/__init__.py
similarity index 100%
rename from adam/jax/__init__.py
rename to src/adam/jax/__init__.py
diff --git a/src/adam/jax/computations.py b/src/adam/jax/computations.py
new file mode 100644
index 00000000..f5174ad0
--- /dev/null
+++ b/src/adam/jax/computations.py
@@ -0,0 +1,204 @@
+# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
+# This software may be modified and distributed under the terms of the
+# GNU Lesser General Public License v2.1 or any later version.
+
+import jax.numpy as jnp
+import numpy as np
+from jax import grad, jit, vmap
+
+from adam.core.rbd_algorithms import RBDAlgorithms
+from adam.jax.jax_like import JaxLike
+
+
+class KinDynComputations(RBDAlgorithms, JaxLike):
+    """This is a small class that retrieves robot quantities using Jax
+    in mixed representation, for Floating Base systems - as humanoid robots.
+    """
+
+    def __init__(
+        self,
+        urdfstring: str,
+        joints_name_list: list,
+        root_link: str = "root_link",
+        gravity: np.array = jnp.array([0, 0, -9.80665, 0, 0, 0]),
+    ) -> None:
+        """
+        Args:
+            urdfstring (str): path of the urdf
+            joints_name_list (list): list of the actuated joints
+            root_link (str, optional): the first link. Defaults to 'root_link'.
+        """
+        super().__init__(
+            urdfstring=urdfstring,
+            joints_name_list=joints_name_list,
+            root_link=root_link,
+            gravity=gravity,
+        )
+
+    def mass_matrix(self, base_transform: jnp.array, joint_positions: jnp.array):
+        """Returns the Mass Matrix functions computed the CRBA
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            M (jax): Mass Matrix
+        """
+        [M, _] = super().crba(base_transform, joint_positions)
+        return M.array
+
+    def centroidal_momentum_matrix(
+        self, base_transform: jnp.array, joint_positions: jnp.array
+    ):
+        """Returns the Centroidal Momentum Matrix functions computed the CRBA
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            Jcc (jnp.array): Centroidal Momentum matrix
+        """
+        [_, Jcm] = self.crba(base_transform, joint_positions)
+        return Jcm.array
+
+    def relative_jacobian(self, frame: str, joint_positions: jnp.array):
+        """Returns the Jacobian between the root link and a specified frame frames
+
+        Args:
+            frame (str): The tip of the chain
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            J (jnp.array): The Jacobian between the root and the frame
+        """
+        return super().relative_jacobian(frame, joint_positions).array
+
+    def forward_kinematics(
+        self, frame: str, base_transform: jnp.array, joint_positions: jnp.array
+    ):
+        """Computes the forward kinematics relative to the specified frame
+
+        Args:
+            frame (str): The frame to which the fk will be computed
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            T_fk (jnp.array): The fk represented as Homogenous transformation matrix
+        """
+        return super().forward_kinematics(frame, base_transform, joint_positions).array
+
+    def forward_kinematics_fun(self, frame):
+        return lambda T, joint_positions: self.forward_kinematics(
+            frame, T, joint_positions
+        )
+
+    def jacobian(self, frame: str, base_transform, joint_positions):
+        """Returns the Jacobian relative to the specified frame
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            s (jnp.array): The joints position
+            frame (str): The frame to which the jacobian will be computed
+
+        Returns:
+            J_tot (jnp.array): The Jacobian relative to the frame
+        """
+        return super().jacobian(frame, base_transform, joint_positions).array
+
+    def bias_force(
+        self,
+        base_transform: jnp.array,
+        joint_positions: jnp.array,
+        base_velocity: jnp.array,
+        s_dot: jnp.array,
+    ) -> jnp.array:
+        """Returns the bias force of the floating-base dynamics ejoint_positionsuation,
+        using a reduced RNEA (no acceleration and external forces)
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+            base_velocity (jnp.array): The base velocity in mixed representation
+            s_dot (jnp.array): The joints velocity
+
+        Returns:
+            h (jnp.array): the bias force
+        """
+        return (
+            super()
+            .rnea(base_transform, joint_positions, base_velocity, s_dot, self.g)
+            .array.squeeze()
+        )
+
+    def coriolis_term(
+        self,
+        base_transform: jnp.array,
+        joint_positions: jnp.array,
+        base_velocity: jnp.array,
+        s_dot: jnp.array,
+    ) -> jnp.array:
+        """Returns the coriolis term of the floating-base dynamics ejoint_positionsuation,
+        using a reduced RNEA (no acceleration and external forces)
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+            base_velocity (jnp.array): The base velocity in mixed representation
+            s_dot (jnp.array): The joints velocity
+
+        Returns:
+            C (jnp.array): the Coriolis term
+        """
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                base_velocity.reshape(6, 1),
+                s_dot,
+                np.zeros(6),
+            )
+            .array.squeeze()
+        )
+
+    def gravity_term(
+        self, base_transform: jnp.array, joint_positions: jnp.array
+    ) -> jnp.array:
+        """Returns the gravity term of the floating-base dynamics ejoint_positionsuation,
+        using a reduced RNEA (no acceleration and external forces)
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            G (jnp.array): the gravity term
+        """
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                np.zeros(6).reshape(6, 1),
+                np.zeros(self.NDoF),
+                self.g,
+            )
+            .array.squeeze()
+        )
+
+    def CoM_position(
+        self, base_transform: jnp.array, joint_positions: jnp.array
+    ) -> jnp.array:
+        """Returns the CoM positon
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            com (jnp.array): The CoM position
+        """
+        return super().CoM_position(base_transform, joint_positions).array.squeeze()
diff --git a/src/adam/jax/jax_like.py b/src/adam/jax/jax_like.py
new file mode 100644
index 00000000..f6cd2fd6
--- /dev/null
+++ b/src/adam/jax/jax_like.py
@@ -0,0 +1,192 @@
+# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
+# This software may be modified and distributed under the terms of the
+# GNU Lesser General Public License v2.1 or any later version.
+
+from dataclasses import dataclass
+from typing import Union
+
+import jax.numpy as jnp
+import numpy.typing as npt
+
+from adam.core.spatial_math import ArrayLike
+
+
+@dataclass
+class JaxLike(ArrayLike):
+    """Wrapper class for Jax types"""
+
+    array: jnp.array
+
+    def __setitem__(self, idx, value: Union["JaxLike", npt.ArrayLike]):
+        """Overrides set item operator"""
+        if type(self) is type(value):
+            value.array = jnp.squeeze(value.array)
+            try:
+                self.array = self.array.at[idx].set(value.array)
+            except:
+                self.array = self.array.at[idx].set(value.array.reshape(-1, 1))
+        else:
+            self.array = self.array.at[idx].set(value)
+
+    def __getitem__(self, idx) -> "JaxLike":
+        """Overrides get item operator"""
+        return JaxLike(self.array[idx])
+
+    @property
+    def shape(self):
+        return self.array.shape
+
+    def reshape(self, *args):
+        return self.array.reshape(*args)
+
+    @property
+    def T(self) -> "JaxLike":
+        """
+        Returns:
+            JaxLike: transpose of the array
+        """
+        return JaxLike(self.array.T)
+
+    def __matmul__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides @ operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array @ jnp.array(other))
+        return JaxLike(self.array @ other.array)
+
+    def __rmatmul__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides @ operator"""
+        if type(self) is not type(other):
+            return JaxLike(jnp.array(other) @ self.array)
+        return JaxLike(other.array @ self.array)
+
+    def __mul__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides * operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array * other)
+        return JaxLike(self.array * other.array)
+
+    def __rmul__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides * operator"""
+        if type(self) is not type(other):
+            return JaxLike(other * self.array)
+        return JaxLike(other.array * self.array)
+
+    def __truediv__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides / operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array / other)
+        return JaxLike(self.array / other.array)
+
+    def __add__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array.squeeze() + other.squeeze())
+        return JaxLike(self.array.squeeze() + other.array.squeeze())
+
+    def __radd__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array.squeeze() + other.squeeze())
+        return JaxLike(self.array.squeeze() + other.array.squeeze())
+
+    def __sub__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides - operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array.squeeze() - other.squeeze())
+        return JaxLike(self.array.squeeze() - other.array.squeeze())
+
+    def __rsub__(self, other: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """Overrides - operator"""
+        if type(self) is not type(other):
+            return JaxLike(self.array.squeeze() - other.squeeze())
+        return JaxLike(self.array.squeeze() - other.array.squeeze())
+
+    def __neg__(self) -> "JaxLike":
+        """Overrides - operator"""
+        return JaxLike(-self.array)
+
+    @staticmethod
+    def zeros(*x) -> "JaxLike":
+        """
+        Returns:
+            JaxLike: Matrix of zeros of dim *x
+        """
+        return JaxLike(jnp.zeros(x))
+
+    @staticmethod
+    def vertcat(*x) -> "JaxLike":
+        """
+        Returns:
+            JaxLike: Vertical concatenation of elements
+        """
+        if isinstance(x[0], JaxLike):
+            v = jnp.vstack([x[i].array for i in range(len(x))]).reshape(-1, 1)
+        else:
+            v = jnp.vstack([x[i] for i in range(len(x))]).reshape(-1, 1)
+        return JaxLike(v)
+
+    @staticmethod
+    def eye(x) -> "JaxLike":
+        """
+        Returns:
+            JaxLike: Identity matrix of dimension x
+        """
+        return JaxLike(jnp.eye(x))
+
+    @staticmethod
+    def skew(x: Union["JaxLike", npt.ArrayLike]) -> "JaxLike":
+        """
+        Args:
+            x (Union[JaxLike, npt.ArrayLike]): vector
+
+        Returns:
+            JaxLike: the skew symmetric matrix from x
+        """
+        if not isinstance(x, JaxLike):
+            return -jnp.cross(jnp.array(x), jnp.eye(3), axisa=0, axisb=0)
+        x = x.array
+        return JaxLike(-jnp.cross(jnp.array(x), jnp.eye(3), axisa=0, axisb=0))
+
+    @staticmethod
+    def array(*x) -> "JaxLike":
+        """
+        Returns:
+            JaxLike: Vector wrapping *x
+        """
+        return JaxLike(jnp.array(x))
+
+    @staticmethod
+    def sin(x: npt.ArrayLike) -> "JaxLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            JaxLike: sin of x
+        """
+        return JaxLike(jnp.sin(x))
+
+    @staticmethod
+    def cos(x: npt.ArrayLike) -> "JaxLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            JaxLike: cos of x
+        """
+        return JaxLike(jnp.cos(x))
+
+    @staticmethod
+    def outer(x: npt.ArrayLike, y: npt.ArrayLike) -> "JaxLike":
+        """
+        Args:
+            x (npt.ArrayLike): vector
+            y (npt.ArrayLike): vector
+
+        Returns:
+            JaxLike: outer product between x and y
+        """
+        x = jnp.array(x)
+        y = jnp.array(y)
+        return JaxLike(jnp.outer(x, y))
diff --git a/adam/numpy/__init__.py b/src/adam/numpy/__init__.py
similarity index 100%
rename from adam/numpy/__init__.py
rename to src/adam/numpy/__init__.py
diff --git a/adam/numpy/computations.py b/src/adam/numpy/computations.py
similarity index 68%
rename from adam/numpy/computations.py
rename to src/adam/numpy/computations.py
index 81c7a9f2..2a951cf5 100644
--- a/adam/numpy/computations.py
+++ b/src/adam/numpy/computations.py
@@ -5,10 +5,10 @@
 import numpy as np
 
 from adam.core.rbd_algorithms import RBDAlgorithms
-from adam.numpy.spatial_math_numpy import SpatialMathNumpy
+from adam.numpy.numpy_like import NumpyLike
 
 
-class KinDynComputations(RBDAlgorithms, SpatialMathNumpy):
+class KinDynComputations(RBDAlgorithms, NumpyLike):
     """This is a small class that retrieves robot quantities using NumPy
     in mixed representation, for Floating Base systems - as humanoid robots.
     """
@@ -33,7 +33,9 @@ def __init__(
             gravity=gravity,
         )
 
-    def mass_matrix(self, base_transform, joint_positions):
+    def mass_matrix(
+        self, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
         """Returns the Mass Matrix functions computed the CRBA
 
         Args:
@@ -44,9 +46,11 @@ def mass_matrix(self, base_transform, joint_positions):
             M (np.ndarray): Mass Matrix
         """
         [M, _] = super().crba(base_transform, joint_positions)
-        return M
+        return M.array
 
-    def centroidal_momentum_matrix(self, base_transform, s):
+    def centroidal_momentum_matrix(
+        self, base_transform: np.ndarray, s: np.ndarray
+    ) -> np.ndarray:
         """Returns the Centroidal Momentum Matrix functions computed the CRBA
 
         Args:
@@ -57,9 +61,11 @@ def centroidal_momentum_matrix(self, base_transform, s):
             Jcc (np.ndarray): Centroidal Momentum matrix
         """
         [_, Jcm] = super().crba(base_transform, s)
-        return Jcm
+        return Jcm.array
 
-    def forward_kinematics(self, frame, base_transform, joint_positions):
+    def forward_kinematics(
+        self, frame: str, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
         """Computes the forward kinematics relative to the specified frame
 
         Args:
@@ -70,9 +76,15 @@ def forward_kinematics(self, frame, base_transform, joint_positions):
         Returns:
             T_fk (np.ndarray): The fk represented as Homogenous transformation matrix
         """
-        return super().forward_kinematics(frame, base_transform, joint_positions)
+        return (
+            super()
+            .forward_kinematics(frame, base_transform, joint_positions)
+            .array.squeeze()
+        )
 
-    def jacobian(self, frame, base_transform, joint_positions):
+    def jacobian(
+        self, frame: str, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
         """Returns the Jacobian relative to the specified frame
 
         Args:
@@ -83,9 +95,9 @@ def jacobian(self, frame, base_transform, joint_positions):
         Returns:
             J_tot (np.ndarray): The Jacobian relative to the frame
         """
-        return super().jacobian(frame, base_transform, joint_positions)
+        return super().jacobian(frame, base_transform, joint_positions).array.squeeze()
 
-    def relative_jacobian(self, frame, joint_positions):
+    def relative_jacobian(self, frame: str, joint_positions: np.ndarray) -> np.ndarray:
         """Returns the Jacobian between the root link and a specified frame frames
 
         Args:
@@ -95,9 +107,11 @@ def relative_jacobian(self, frame, joint_positions):
         Returns:
             J (np.ndarray): The Jacobian between the root and the frame
         """
-        return super().relative_jacobian(frame, joint_positions)
+        return super().relative_jacobian(frame, joint_positions).array
 
-    def CoM_position(self, base_transform, joint_positions):
+    def CoM_position(
+        self, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
         """Returns the CoM positon
 
         Args:
@@ -105,13 +119,17 @@ def CoM_position(self, base_transform, joint_positions):
             joint_positions (np.ndarray): The joints position
 
         Returns:
-            com (np.ndarray): The CoM position
+            CoM (np.ndarray): The CoM position
         """
-        return super().CoM_position(base_transform, joint_positions)
+        return super().CoM_position(base_transform, joint_positions).array.squeeze()
 
     def bias_force(
-        self, base_transform, joint_positions, base_velocity, joint_velocities
-    ):
+        self,
+        base_transform: np.ndarray,
+        joint_positions: np.ndarray,
+        base_velocity: np.ndarray,
+        joint_velocities: np.ndarray,
+    ) -> np.ndarray:
         """Returns the bias force of the floating-base dynamics equation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -124,18 +142,25 @@ def bias_force(
         Returns:
             h (np.ndarray): the bias force
         """
-        h = super().rnea(
-            base_transform,
-            joint_positions,
-            base_velocity.reshape(6, 1),
-            joint_velocities,
-            self.g,
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                base_velocity.reshape(6, 1),
+                joint_velocities,
+                self.g,
+            )
+            .array.squeeze()
         )
-        return h[:, 0]
 
     def coriolis_term(
-        self, base_transform, joint_positions, base_velocity, joint_velocities
-    ):
+        self,
+        base_transform: np.ndarray,
+        joint_positions: np.ndarray,
+        base_velocity: np.ndarray,
+        joint_velocities: np.ndarray,
+    ) -> np.ndarray:
         """Returns the coriolis term of the floating-base dynamics equation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -149,16 +174,21 @@ def coriolis_term(
             C (np.ndarray): the Coriolis term
         """
         # set in the bias force computation the gravity term to zero
-        C = super().rnea(
-            base_transform,
-            joint_positions,
-            base_velocity.reshape(6, 1),
-            joint_velocities,
-            np.zeros(6),
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                base_velocity.reshape(6, 1),
+                joint_velocities,
+                np.zeros(6),
+            )
+            .array.squeeze()
         )
-        return C[:, 0]
 
-    def gravity_term(self, base_transform, joint_positions):
+    def gravity_term(
+        self, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
         """Returns the gravity term of the floating-base dynamics equation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -169,11 +199,14 @@ def gravity_term(self, base_transform, joint_positions):
         Returns:
             G (np.ndarray): the gravity term
         """
-        G = super().rnea(
-            base_transform,
-            joint_positions,
-            np.zeros(6).reshape(6, 1),
-            np.zeros(self.NDoF),
-            self.g,
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                np.zeros(6).reshape(6, 1),
+                np.zeros(self.NDoF),
+                self.g,
+            )
+            .array.squeeze()
         )
-        return G[:, 0]
diff --git a/src/adam/numpy/numpy_like.py b/src/adam/numpy/numpy_like.py
new file mode 100644
index 00000000..4cfdcd46
--- /dev/null
+++ b/src/adam/numpy/numpy_like.py
@@ -0,0 +1,200 @@
+# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
+# This software may be modified and distributed under the terms of the
+# GNU Lesser General Public License v2.1 or any later version.
+
+from dataclasses import dataclass
+from typing import Union
+
+import numpy as np
+import numpy.typing as npt
+
+from adam.core.spatial_math import ArrayLike
+
+
+@dataclass
+class NumpyLike(ArrayLike):
+    """Class wrapping NumPy types"""
+
+    array: np.ndarray
+
+    def __setitem__(self, idx, value: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides set item operator"""
+        if type(self) is type(value):
+            value.array = np.squeeze(value.array)
+            try:
+                self.array[idx] = value.array
+            except:
+                self.array[idx] = value.array.reshape(-1, 1)
+        else:
+            self.array[idx] = value
+
+    def __getitem__(self, idx) -> "NumpyLike":
+        """Overrides get item operator"""
+        return NumpyLike(self.array[idx])
+
+    @property
+    def shape(self):
+        return self.array.shape
+
+    def reshape(self, *args):
+        return self.array.reshape(*args)
+
+    @property
+    def T(self) -> "NumpyLike":
+        """
+        Returns:
+            NumpyLike: transpose of the array
+        """
+        return NumpyLike(self.array.T)
+
+    def __matmul__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return NumpyLike(self.array @ other.array)
+        else:
+            return NumpyLike(self.array @ np.array(other))
+
+    def __rmatmul__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return NumpyLike(other.array @ self.array)
+        else:
+            return NumpyLike(other @ self.array)
+
+    def __mul__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return NumpyLike(self.array * other.array)
+        else:
+            return NumpyLike(self.array * other)
+
+    def __rmul__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return NumpyLike(other.array * self.array)
+        else:
+            return NumpyLike(other * self.array)
+
+    def __truediv__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides / operator"""
+        if type(self) is type(other):
+            return NumpyLike(self.array / other.array)
+        else:
+            return NumpyLike(self.array / other)
+
+    def __add__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return NumpyLike(self.array.squeeze() + other.squeeze())
+        return NumpyLike(self.array.squeeze() + other.array.squeeze())
+
+    def __radd__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return NumpyLike(self.array + other)
+        return NumpyLike(self.array + other.array)
+
+    def __sub__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides - operator"""
+        if type(self) is not type(other):
+            return NumpyLike(self.array.squeeze() - other.squeeze())
+        return NumpyLike(self.array.squeeze() - other.array.squeeze())
+
+    def __rsub__(self, other: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """Overrides - operator"""
+        if type(self) is not type(other):
+            return NumpyLike(self.array.squeeze() - other.squeeze())
+        return NumpyLike(self.array.squeeze() - other.array.squeeze())
+
+    def __neg__(self):
+        """Overrides - operator"""
+        return NumpyLike(-self.array)
+
+    @staticmethod
+    def zeros(*x) -> "NumpyLike":
+        """
+        Returns:
+            NumpyLike: zero matrix of dimension x
+        """
+        return NumpyLike(np.zeros(x))
+
+    @staticmethod
+    def vertcat(*x: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """
+        Returns:
+            NumpyLike: vertical concatenation of x
+        """
+        if isinstance(x[0], NumpyLike):
+            v = np.vstack([x[i].array for i in range(len(x))]).reshape(-1, 1)
+        else:
+            v = np.vstack([x[i] for i in range(len(x))]).reshape(-1, 1)
+        return NumpyLike(v)
+
+    @staticmethod
+    def eye(x: int) -> "NumpyLike":
+        """
+        Args:
+            x (int): matrix dimension
+
+        Returns:
+            NumpyLike: Identity matrix of dimension x
+        """
+        return NumpyLike(np.eye(x))
+
+    @staticmethod
+    def skew(x: Union["NumpyLike", npt.ArrayLike]) -> "NumpyLike":
+        """
+        Args:
+            x (Union[NumpyLike, npt.ArrayLike]): vector
+
+        Returns:
+            NumpyLike:  the skew symmetric matrix from x
+        """
+        if not isinstance(x, NumpyLike):
+            return -np.cross(np.array(x), np.eye(3), axisa=0, axisb=0)
+        x = x.array
+        return NumpyLike(-np.cross(np.array(x), np.eye(3), axisa=0, axisb=0))
+
+    @staticmethod
+    def array(*x) -> "NumpyLike":
+        """
+        Returns:
+            NumpyLike: Vector wrapping *x
+        """
+        return NumpyLike(np.array(x))
+
+    @staticmethod
+    def sin(x: npt.ArrayLike) -> "NumpyLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            NumpyLike: sin value of x
+        """
+        return NumpyLike(np.sin(x))
+
+    @staticmethod
+    def cos(x: npt.ArrayLike) -> "NumpyLike":
+        """
+        Args:
+            x (npt.ArrayLike): angle value
+
+        Returns:
+            NumpyLike: cos value of x
+        """
+        return NumpyLike(np.cos(x))
+
+    @staticmethod
+    def outer(x: npt.ArrayLike, y: npt.ArrayLike) -> "NumpyLike":
+        """
+        Args:
+            x (npt.ArrayLike): vector
+            y (npt.ArrayLike): vector
+
+        Returns:
+            NumpyLike: outer product of x and y
+        """
+        x = np.array(x)
+        y = np.array(y)
+        return NumpyLike(np.outer(x, y))
diff --git a/adam/pytorch/__init__.py b/src/adam/pytorch/__init__.py
similarity index 100%
rename from adam/pytorch/__init__.py
rename to src/adam/pytorch/__init__.py
diff --git a/adam/pytorch/computations.py b/src/adam/pytorch/computations.py
similarity index 51%
rename from adam/pytorch/computations.py
rename to src/adam/pytorch/computations.py
index 7c53aa12..1e4f29b3 100644
--- a/adam/pytorch/computations.py
+++ b/src/adam/pytorch/computations.py
@@ -6,10 +6,10 @@
 import torch
 
 from adam.core.rbd_algorithms import RBDAlgorithms
-from adam.pytorch.spatial_math_pytorch import SpatialMathPytorch
+from adam.pytorch.torch_like import TorchLike
 
 
-class KinDynComputations(RBDAlgorithms, SpatialMathPytorch):
+class KinDynComputations(RBDAlgorithms, TorchLike):
     """This is a small class that retrieves robot quantities using Pytorch
     in mixed representation, for Floating Base systems - as humanoid robots.
     """
@@ -34,7 +34,9 @@ def __init__(
             gravity=gravity,
         )
 
-    def mass_matrix(self, base_transform, s):
+    def mass_matrix(
+        self, base_transform: torch.Tensor, s: torch.Tensor
+    ) -> torch.Tensor:
         """Returns the Mass Matrix functions computed the CRBA
 
         Args:
@@ -45,9 +47,11 @@ def mass_matrix(self, base_transform, s):
             M (torch.tensor): Mass Matrix
         """
         [M, _] = super().crba(base_transform, s)
-        return M
+        return M.array
 
-    def centroidal_momentum_matrix(self, base_transform, s):
+    def centroidal_momentum_matrix(
+        self, base_transform: torch.Tensor, s: torch.Tensor
+    ) -> torch.Tensor:
         """Returns the Centroidal Momentum Matrix functions computed the CRBA
 
         Args:
@@ -58,9 +62,11 @@ def centroidal_momentum_matrix(self, base_transform, s):
             Jcc (torch.tensor): Centroidal Momentum matrix
         """
         [_, Jcm] = super().crba(base_transform, s)
-        return Jcm
+        return Jcm.array
 
-    def forward_kinematics(self, frame, base_transform, s):
+    def forward_kinematics(
+        self, frame, base_transform: torch.Tensor, s: torch.Tensor
+    ) -> torch.Tensor:
         """Computes the forward kinematics relative to the specified frame
 
         Args:
@@ -71,11 +77,15 @@ def forward_kinematics(self, frame, base_transform, s):
         Returns:
             T_fk (torch.tensor): The fk represented as Homogenous transformation matrix
         """
-        return super().forward_kinematics(
-            frame, torch.FloatTensor(base_transform), torch.FloatTensor(s)
-        )
-
-    def jacobian(self, frame, base_transform, joint_positions):
+        return (
+            super().forward_kinematics(
+                frame, torch.FloatTensor(base_transform), torch.FloatTensor(s)
+            )
+        ).array
+
+    def jacobian(
+        self, frame: str, base_transform: torch.Tensor, joint_positions: torch.Tensor
+    ) -> torch.Tensor:
         """Returns the Jacobian relative to the specified frame
 
         Args:
@@ -86,48 +96,9 @@ def jacobian(self, frame, base_transform, joint_positions):
         Returns:
             J_tot (torch.tensor): The Jacobian relative to the frame
         """
-        chain = self.robot_desc.get_chain(self.root_link, frame)
-        T_fk = self.eye(4)
-        T_fk = T_fk @ base_transform
-        J = self.zeros(6, self.NDoF)
-        T_ee = self.forward_kinematics(frame, base_transform, joint_positions)
-        P_ee = T_ee[:3, 3]
-        for item in chain:
-            if item in self.robot_desc.joint_map:
-                joint = self.robot_desc.joint_map[item]
-                if joint.type == "fixed":
-                    xyz = joint.origin.xyz
-                    rpy = joint.origin.rpy
-                    joint_frame = self.H_from_Pos_RPY(xyz, rpy)
-                    T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
-                    T_joint = self.H_revolute_joint(
-                        joint.origin.xyz,
-                        joint.origin.rpy,
-                        joint.axis,
-                        q_,
-                    )
-                    T_fk = T_fk @ T_joint
-                    p_prev = P_ee - T_fk[:3, 3]
-                    z_prev = T_fk[:3, :3] @ torch.tensor(joint.axis)
-                    if joint.idx is not None:
-                        stack = np.hstack(((self.skew(z_prev) @ p_prev), z_prev))
-                        J[:, joint.idx] = torch.tensor(stack)
-
-        # Adding the floating base part of the Jacobian, in Mixed representation
-        J_tot = self.zeros(6, self.NDoF + 6)
-        J_tot[:3, :3] = self.eye(3)
-        J_tot[:3, 3:6] = -self.skew((P_ee - base_transform[:3, 3]))
-        J_tot[:3, 6:] = J[:3, :]
-        J_tot[3:, 3:6] = self.eye(3)
-        J_tot[3:, 6:] = J[3:, :]
-        return J_tot
-
-    def relative_jacobian(self, frame, joint_positions):
+        return super().jacobian(frame, base_transform, joint_positions).array
+
+    def relative_jacobian(self, frame, joint_positions: torch.Tensor) -> torch.Tensor:
         """Returns the Jacobian between the root link and a specified frame frames
 
         Args:
@@ -137,41 +108,11 @@ def relative_jacobian(self, frame, joint_positions):
         Returns:
             J (torch.tensor): The Jacobian between the root and the frame
         """
-        chain = self.robot_desc.get_chain(self.root_link, frame)
-        base_transform = self.eye(4)
-        T_fk = self.eye(4)
-        T_fk = T_fk @ base_transform
-        J = self.zeros(6, self.NDoF)
-        T_ee = self.forward_kinematics(frame, base_transform, joint_positions)
-        P_ee = T_ee[:3, 3]
-        for item in chain:
-            if item in self.robot_desc.joint_map:
-                joint = self.robot_desc.joint_map[item]
-                if joint.type == "fixed":
-                    xyz = joint.origin.xyz
-                    rpy = joint.origin.rpy
-                    joint_frame = self.H_from_Pos_RPY(xyz, rpy)
-                    T_fk = T_fk @ joint_frame
-                if joint.type == "revolute" or joint.type == "continuous":
-                    if joint.idx is not None:
-                        q_ = joint_positions[joint.idx]
-                    else:
-                        q_ = 0.0
-                    T_joint = self.H_revolute_joint(
-                        joint.origin.xyz,
-                        joint.origin.rpy,
-                        joint.axis,
-                        q_,
-                    )
-                    T_fk = T_fk @ T_joint
-                    p_prev = P_ee - T_fk[:3, 3]
-                    z_prev = T_fk[:3, :3] @ torch.tensor(joint.axis)
-                    if joint.idx is not None:
-                        stack = np.hstack(((self.skew(z_prev) @ p_prev), z_prev))
-                        J[:, joint.idx] = torch.tensor(stack)
-        return J
-
-    def CoM_position(self, base_transform, joint_positions):
+        return super().relative_jacobian(frame, joint_positions).array
+
+    def CoM_position(
+        self, base_transform: torch.Tensor, joint_positions: torch.Tensor
+    ) -> torch.Tensor:
         """Returns the CoM positon
 
         Args:
@@ -181,9 +122,15 @@ def CoM_position(self, base_transform, joint_positions):
         Returns:
             com (torch.tensor): The CoM position
         """
-        return super().CoM_position(base_transform, joint_positions)
+        return super().CoM_position(base_transform, joint_positions).array.squeeze()
 
-    def bias_force(self, base_transform, s, base_velocity, joint_velocities):
+    def bias_force(
+        self,
+        base_transform: torch.Tensor,
+        s: torch.Tensor,
+        base_velocity: torch.Tensor,
+        joint_velocities: torch.Tensor,
+    ) -> torch.Tensor:
         """Returns the bias force of the floating-base dynamics ejoint_positionsuation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -196,14 +143,25 @@ def bias_force(self, base_transform, s, base_velocity, joint_velocities):
         Returns:
             h (torch.tensor): the bias force
         """
-        h = super().rnea(
-            base_transform, s, base_velocity.reshape(6, 1), joint_velocities, self.g
+        return (
+            super()
+            .rnea(
+                base_transform,
+                s,
+                base_velocity.reshape(6, 1),
+                joint_velocities,
+                self.g,
+            )
+            .array.squeeze()
         )
-        return h[:, 0]
 
     def coriolis_term(
-        self, base_transform, joint_positions, base_velocity, joint_velocities
-    ):
+        self,
+        base_transform: torch.Tensor,
+        joint_positions: torch.Tensor,
+        base_velocity: torch.Tensor,
+        joint_velocities: torch.Tensor,
+    ) -> torch.Tensor:
         """Returns the coriolis term of the floating-base dynamics ejoint_positionsuation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -217,16 +175,21 @@ def coriolis_term(
             C (torch.tensor): the Coriolis term
         """
         # set in the bias force computation the gravity term to zero
-        C = super().rnea(
-            base_transform,
-            joint_positions,
-            base_velocity.reshape(6, 1),
-            joint_velocities,
-            torch.zeros(6),
+        return (
+            super()
+            .rnea(
+                base_transform,
+                joint_positions,
+                base_velocity.reshape(6, 1),
+                joint_velocities,
+                torch.zeros(6),
+            )
+            .array.squeeze()
         )
-        return C[:, 0]
 
-    def gravity_term(self, base_transform, base_positions):
+    def gravity_term(
+        self, base_transform: torch.Tensor, base_positions: torch.Tensor
+    ) -> torch.Tensor:
         """Returns the gravity term of the floating-base dynamics ejoint_positionsuation,
         using a reduced RNEA (no acceleration and external forces)
 
@@ -237,11 +200,14 @@ def gravity_term(self, base_transform, base_positions):
         Returns:
             G (torch.tensor): the gravity term
         """
-        G = super().rnea(
-            base_transform,
-            base_positions,
-            torch.zeros(6).reshape(6, 1),
-            torch.zeros(self.NDoF),
-            torch.FloatTensor(self.g),
+        return (
+            super()
+            .rnea(
+                base_transform,
+                base_positions,
+                torch.zeros(6).reshape(6, 1),
+                torch.zeros(self.NDoF),
+                torch.FloatTensor(self.g),
+            )
+            .array.squeeze()
         )
-        return G[:, 0]
diff --git a/src/adam/pytorch/torch_like.py b/src/adam/pytorch/torch_like.py
new file mode 100644
index 00000000..ca47515e
--- /dev/null
+++ b/src/adam/pytorch/torch_like.py
@@ -0,0 +1,208 @@
+# Copyright (C) 2021 Istituto Italiano di Tecnologia (IIT). All rights reserved.
+# This software may be modified and distributed under the terms of the
+# GNU Lesser General Public License v2.1 or any later version.
+
+from dataclasses import dataclass
+from typing import Union
+
+import numpy.typing as ntp
+import torch
+
+from adam.core.spatial_math import ArrayLike
+
+
+@dataclass
+class TorchLike(ArrayLike):
+    """Class wrapping pyTorch types"""
+
+    array: torch.Tensor
+
+    def __setitem__(self, idx, value: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides set item operator"""
+        if type(self) is type(value):
+            value.array = torch.squeeze(value.array)
+            try:
+                self.array[idx] = value.array
+            except:
+                self.array[idx] = value.array.reshape(-1, 1)
+        else:
+            self.array[idx] = torch.FloatTensor(value)
+
+    def __getitem__(self, idx):
+        """Overrides get item operator"""
+        return TorchLike(self.array[idx])
+
+    @property
+    def shape(self):
+        return self.array.shape
+
+    def reshape(self, *args):
+        return self.array.reshape(*args)
+
+    @property
+    def T(self) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: transpose of array
+        """
+        return TorchLike(self.array.T)
+
+    def __matmul__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return TorchLike(self.array @ other.array)
+        else:
+            return TorchLike(self.array @ torch.FloatTensor(other))
+
+    def __rmatmul__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides @ operator"""
+        if type(self) is type(other):
+            return TorchLike(other.array @ self.array)
+        else:
+            return TorchLike(torch.FloatTensor(other) @ self.array)
+
+    def __mul__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return TorchLike(self.array * other.array)
+        else:
+            return TorchLike(self.array * other)
+
+    def __rmul__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides * operator"""
+        if type(self) is type(other):
+            return TorchLike(other.array * self.array)
+        else:
+            return TorchLike(other * self.array)
+
+    def __truediv__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides / operator"""
+        if type(self) is type(other):
+            return TorchLike(self.array / other.array)
+        else:
+            return TorchLike(self.array / other)
+
+    def __add__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return TorchLike(self.array.squeeze() + other.squeeze())
+        return TorchLike(self.array.squeeze() + other.array.squeeze())
+
+    def __radd__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides + operator"""
+        if type(self) is not type(other):
+            return TorchLike(self.array.squeeze() + other.squeeze())
+        return TorchLike(self.array.squeeze() + other.array.squeeze())
+
+    def __sub__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides - operator"""
+        if type(self) is type(other):
+            return TorchLike(self.array.squeeze() - other.array.squeeze())
+        else:
+            return TorchLike(self.array.squeeze() - other.squeeze())
+
+    def __rsub__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """Overrides - operator"""
+        if type(self) is type(other):
+            return TorchLike(other.array.squeeze() - self.array.squeeze())
+        else:
+            return TorchLike(other.squeeze() - self.array.squeeze())
+
+    def __neg__(self) -> "TorchLike":
+        """Overrides - operator"""
+        return TorchLike(-self.array)
+
+    @staticmethod
+    def zeros(*x: int) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: zero matrix of dimension *x
+        """
+        return TorchLike(torch.zeros(x).float())
+
+    @staticmethod
+    def vertcat(*x: ntp.ArrayLike) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: vertical concatenation of x
+        """
+        if isinstance(x[0], TorchLike):
+            v = torch.vstack([x[i].array for i in range(len(x))]).reshape(-1, 1)
+        else:
+            v = torch.FloatTensor(x).reshape(-1, 1)
+        return TorchLike(v)
+
+    @staticmethod
+    def eye(x: int) -> "TorchLike":
+        """
+        Args:
+            x (int): dimension
+
+        Returns:
+            TorchLike: identity matrix of dimension x
+        """
+        return TorchLike(torch.eye(x).float())
+
+    @staticmethod
+    def skew(x: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """
+        Args:
+            x (Union[TorchLike, ntp.ArrayLike]): vector
+
+        Returns:
+            TorchLike: skew matrix from x
+        """
+        if not isinstance(x, TorchLike):
+            return TorchLike(
+                torch.FloatTensor(
+                    [[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]]
+                )
+            )
+        x = x.array
+        return TorchLike(
+            torch.FloatTensor([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
+        )
+
+    @staticmethod
+    def array(*x: ntp.ArrayLike) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: vector wrapping x
+        """
+        return TorchLike(torch.FloatTensor(x))
+
+    @staticmethod
+    def sin(x: ntp.ArrayLike) -> "TorchLike":
+        """
+        Args:
+            x (ntp.ArrayLike): angle value
+
+        Returns:
+            TorchLike: sin value of x
+        """
+        x = torch.tensor(x)
+        return TorchLike(torch.sin(x))
+
+    @staticmethod
+    def cos(x: ntp.ArrayLike) -> "TorchLike":
+        """
+        Args:
+            x (ntp.ArrayLike): angle value
+
+        Returns:
+            TorchLike: cos value of x
+        """
+        x = torch.tensor(x)
+        return TorchLike(torch.cos(x))
+
+    @staticmethod
+    def outer(x: ntp.ArrayLike, y: ntp.ArrayLike) -> "TorchLike":
+        """
+        Args:
+            x (ntp.ArrayLike): vector
+            y (ntp.ArrayLike): vector
+
+        Returns:
+            TorchLike: outer product of x and y
+        """
+        return TorchLike(torch.outer(torch.tensor(x), torch.tensor(y)))


Test patch:
diff --git a/tests/test_CasADi_computations.py b/tests/test_CasADi_computations.py
index fb7746f4..d19ccd88 100644
--- a/tests/test_CasADi_computations.py
+++ b/tests/test_CasADi_computations.py
@@ -42,10 +42,6 @@
 ]
 
 
-def SX2DM(x):
-    return cs.DM(x)
-
-
 def H_from_Pos_RPY_idyn(xyz, rpy):
     T = idyntree.Transform.Identity()
     R = idyntree.Rotation.RPY(rpy[0], rpy[1], rpy[2])
@@ -104,7 +100,7 @@ def test_mass_matrix():
     mass_mx = idyntree.MatrixDynSize()
     kinDyn.getFreeFloatingMassMatrix(mass_mx)
     mass_mxNumpy = mass_mx.toNumPy()
-    mass_test = SX2DM(M(H_b, s_))
+    mass_test = cs.DM(M(H_b, s_))
     assert mass_test - mass_mxNumpy == pytest.approx(0.0, abs=1e-5)
 
 
@@ -113,13 +109,13 @@ def test_CMM():
     cmm_idyntree = idyntree.MatrixDynSize()
     kinDyn.getCentroidalTotalMomentumJacobian(cmm_idyntree)
     cmm_idyntreeNumpy = cmm_idyntree.toNumPy()
-    Jcm_test = SX2DM(Jcm(H_b, s_))
+    Jcm_test = cs.DM(Jcm(H_b, s_))
     assert Jcm_test - cmm_idyntreeNumpy == pytest.approx(0.0, abs=1e-5)
 
 
 def test_CoM_pos():
     com_f = comp.CoM_position_fun()
-    CoM_cs = SX2DM(com_f(H_b, s_))
+    CoM_cs = cs.DM(com_f(H_b, s_))
     CoM_iDynTree = kinDyn.getCenterOfMassPosition().toNumPy()
     assert CoM_cs - CoM_iDynTree == pytest.approx(0.0, abs=1e-5)
 
@@ -135,7 +131,7 @@ def test_jacobian():
     iDyntreeJ_ = idyntree.MatrixDynSize(6, n_dofs + 6)
     kinDyn.getFrameFreeFloatingJacobian("l_sole", iDyntreeJ_)
     iDynNumpyJ_ = iDyntreeJ_.toNumPy()
-    J_test = SX2DM(J_tot(H_b, s_))
+    J_test = cs.DM(J_tot(H_b, s_))
     assert iDynNumpyJ_ - J_test == pytest.approx(0.0, abs=1e-5)
 
 
@@ -144,7 +140,7 @@ def test_jacobian_non_actuated():
     iDyntreeJ_ = idyntree.MatrixDynSize(6, n_dofs + 6)
     kinDyn.getFrameFreeFloatingJacobian("head", iDyntreeJ_)
     iDynNumpyJ_ = iDyntreeJ_.toNumPy()
-    J_test = SX2DM(J_tot(H_b, s_))
+    J_test = cs.DM(J_tot(H_b, s_))
     assert iDynNumpyJ_ - J_test == pytest.approx(0.0, abs=1e-5)
 
 
@@ -153,7 +149,7 @@ def test_fk():
     p_idy2np = H_idyntree.getPosition().toNumPy()
     R_idy2np = H_idyntree.getRotation().toNumPy()
     T = comp.forward_kinematics_fun("l_sole")
-    H_test = SX2DM(T(H_b, s_))
+    H_test = cs.DM(T(H_b, s_))
     assert R_idy2np - H_test[:3, :3] == pytest.approx(0.0, abs=1e-5)
     assert p_idy2np - H_test[:3, 3] == pytest.approx(0.0, abs=1e-5)
 
@@ -163,7 +159,7 @@ def test_fk_non_actuated():
     p_idy2np = H_idyntree.getPosition().toNumPy()
     R_idy2np = H_idyntree.getRotation().toNumPy()
     T = comp.forward_kinematics_fun("head")
-    H_test = SX2DM(T(H_b, s_))
+    H_test = cs.DM(T(H_b, s_))
     assert R_idy2np - H_test[:3, :3] == pytest.approx(0.0, abs=1e-5)
     assert p_idy2np - H_test[:3, 3] == pytest.approx(0.0, abs=1e-5)
 
@@ -175,7 +171,7 @@ def test_bias_force():
         (h_iDyn.baseWrench().toNumPy(), h_iDyn.jointTorques().toNumPy())
     )
     h = comp.bias_force_fun()
-    h_test = SX2DM(h(H_b, s_, vb_, s_dot_))
+    h_test = cs.DM(h(H_b, s_, vb_, s_dot_))
     assert h_iDyn_np - h_test == pytest.approx(0.0, abs=1e-4)
 
 
@@ -189,7 +185,7 @@ def test_coriolis_term():
         (C_iDyn.baseWrench().toNumPy(), C_iDyn.jointTorques().toNumPy())
     )
     C = comp.coriolis_term_fun()
-    C_test = SX2DM(C(H_b, s_, vb_, s_dot_))
+    C_test = cs.DM(C(H_b, s_, vb_, s_dot_))
     assert C_iDyn_np - C_test == pytest.approx(0.0, abs=1e-4)
 
 
@@ -203,5 +199,5 @@ def test_gravity_term():
         (G_iDyn.baseWrench().toNumPy(), G_iDyn.jointTorques().toNumPy())
     )
     G = comp.gravity_term_fun()
-    G_test = SX2DM(G(H_b, s_))
+    G_test = cs.DM(G(H_b, s_))
     assert G_iDyn_np - G_test == pytest.approx(0.0, abs=1e-4)
diff --git a/tests/test_pytorch_computations.py b/tests/test_pytorch_computations.py
index 56bc9933..34e0c16f 100644
--- a/tests/test_pytorch_computations.py
+++ b/tests/test_pytorch_computations.py
@@ -100,7 +100,7 @@ def test_mass_matrix():
     kinDyn.getFreeFloatingMassMatrix(mass_mx)
     mass_mxNumpy = mass_mx.toNumPy()
     mass_test = comp.mass_matrix(H_b, s_)
-    assert mass_test - mass_mxNumpy == pytest.approx(0.0, abs=1e-5)
+    assert mass_test - mass_mxNumpy == pytest.approx(0.0, abs=1e-4)
 
 
 def test_CMM():
@@ -108,18 +108,18 @@ def test_CMM():
     kinDyn.getCentroidalTotalMomentumJacobian(cmm_idyntree)
     cmm_idyntreeNumpy = cmm_idyntree.toNumPy()
     Jcm_test = comp.centroidal_momentum_matrix(H_b, s_)
-    assert Jcm_test - cmm_idyntreeNumpy == pytest.approx(0.0, abs=1e-5)
+    assert Jcm_test - cmm_idyntreeNumpy == pytest.approx(0.0, abs=1e-4)
 
 
 def test_CoM_pos():
     CoM_test = comp.CoM_position(H_b, s_)
     CoM_iDynTree = kinDyn.getCenterOfMassPosition().toNumPy()
-    assert CoM_test - CoM_iDynTree == pytest.approx(0.0, abs=1e-5)
+    assert CoM_test - CoM_iDynTree == pytest.approx(0.0, abs=1e-4)
 
 
 def test_total_mass():
     assert comp.get_total_mass() - robot_iDyn.model().getTotalMass() == pytest.approx(
-        0.0, abs=1e-5
+        0.0, abs=1e-4
     )
 
 
@@ -128,7 +128,7 @@ def test_jacobian():
     kinDyn.getFrameFreeFloatingJacobian("l_sole", iDyntreeJ_)
     iDynNumpyJ_ = iDyntreeJ_.toNumPy()
     J_test = comp.jacobian("l_sole", H_b, s_)
-    assert iDynNumpyJ_ - np.array(J_test) == pytest.approx(0.0, abs=1e-5)
+    assert iDynNumpyJ_ - np.array(J_test) == pytest.approx(0.0, abs=1e-4)
 
 
 def test_jacobian_non_actuated():
@@ -136,7 +136,7 @@ def test_jacobian_non_actuated():
     kinDyn.getFrameFreeFloatingJacobian("head", iDyntreeJ_)
     iDynNumpyJ_ = iDyntreeJ_.toNumPy()
     J_test = comp.jacobian("head", H_b, s_)
-    assert iDynNumpyJ_ - np.array(J_test) == pytest.approx(0.0, abs=1e-5)
+    assert iDynNumpyJ_ - np.array(J_test) == pytest.approx(0.0, abs=1e-4)
 
 
 def test_fk():
@@ -144,8 +144,8 @@ def test_fk():
     p_idy2np = H_idyntree.getPosition().toNumPy()
     R_idy2np = H_idyntree.getRotation().toNumPy()
     H_test = comp.forward_kinematics("l_sole", H_b, s_)
-    assert torch.tensor(R_idy2np) - H_test[:3, :3] == pytest.approx(0.0, abs=1e-5)
-    assert torch.tensor(p_idy2np) - H_test[:3, 3] == pytest.approx(0.0, abs=1e-5)
+    assert (R_idy2np) - np.array(H_test[:3, :3]) == pytest.approx(0.0, abs=1e-4)
+    assert (p_idy2np) - np.array(H_test[:3, 3]) == pytest.approx(0.0, abs=1e-4)
 
 
 def test_fk_non_actuated():
@@ -153,8 +153,8 @@ def test_fk_non_actuated():
     p_idy2np = H_idyntree.getPosition().toNumPy()
     R_idy2np = H_idyntree.getRotation().toNumPy()
     H_test = comp.forward_kinematics("head", H_b, s_)
-    assert R_idy2np - np.array(H_test[:3, :3]) == pytest.approx(0.0, abs=1e-5)
-    assert p_idy2np - np.array(H_test[:3, 3]) == pytest.approx(0.0, abs=1e-5)
+    assert R_idy2np - np.array(H_test[:3, :3]) == pytest.approx(0.0, abs=1e-4)
+    assert p_idy2np - np.array(H_test[:3, 3]) == pytest.approx(0.0, abs=1e-4)
 
 
 def test_bias_force():
@@ -203,4 +203,4 @@ def test_relative_jacobian():
     kinDyn.getFrameFreeFloatingJacobian("l_sole", iDyntreeJ_)
     iDynNumpyRelativeJ = (iDyntreeJ_.toNumPy())[:, 6:]
     J_test = comp.relative_jacobian("l_sole", s_)
-    assert iDynNumpyRelativeJ - np.array(J_test) == pytest.approx(0.0, abs=1e-5)
+    assert iDynNumpyRelativeJ - np.array(J_test) == pytest.approx(0.0, abs=1e-4)