Road to BatchedPytorchWithoutJax - PR1

src/adam/casadi/casadi_like.py

@@ -1,7 +1,7 @@
 # Copyright (C) Istituto Italiano di Tecnologia (IIT). All rights reserved.
 
 from dataclasses import dataclass
-from typing import Union
+from typing import Union, Tuple
 
 import casadi as cs
 import numpy.typing as npt
@@ -85,11 +85,66 @@ def __truediv__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike"
 
     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
+        if idx is Ellipsis:
+            self.array = value.array if isinstance(value, CasadiLike) else value
+        elif isinstance(idx, tuple) and Ellipsis in idx:
+            idx = tuple(slice(None) if i is Ellipsis else i for i in idx)
+            self.array[idx] = value.array if isinstance(value, CasadiLike) else value
+        else:
+            self.array[idx] = value.array if isinstance(value, CasadiLike) else value
+
+    @property
+    def shape(self) -> Tuple[int]:
+        """
+        Returns:
+            Tuple[int]: shape of the array
+        """
+
+        # We force to have the same interface as numpy
+        if self.array.shape[1] == 1 and self.array.shape[0] == 1:
+            return tuple()
+        elif self.array.shape[1] == 1:
+            return (self.array.shape[0],)
+
+        return self.array.shape
+
+    def reshape(self, *args) -> "CasadiLike":
+        """
+        Args:
+            *args: new shape
+        """
+        args = tuple(filter(None, args))
+        if len(args) > 2:
+            raise ValueError(
+                f"Only 1D and 2D arrays are supported, The shape is {args}"
+            )
+
+        # For 1D reshape, just call CasADi reshape directly
+        if len(args) == 1:
+            new_array = cs.reshape(self.array, args[0], 1)
+        else:
+            # For 2D reshape, transpose before and after to mimic row-major behavior
+            new_array = cs.reshape(self.array.T, args[1], args[0]).T
+
+        return CasadiLike(new_array)
 
     def __getitem__(self, idx) -> "CasadiLike":
         """Overrides get item operator"""
-        return CasadiLike(self.array[idx])
+        if idx is Ellipsis:
+            # Handle the case where only Ellipsis is passed
+            return CasadiLike(self.array)
+        elif isinstance(idx, tuple) and Ellipsis in idx:
+            if len(self.shape) == 2:
+                idx = tuple(slice(None) if k is Ellipsis else k for k in idx)
+            else:
+                # take the values that are not Ellipsis
+                idx = idx[: idx.index(Ellipsis)] + idx[idx.index(Ellipsis) + 1 :]
+                idx = idx[0] if len(idx) == 1 else idx
+
+            return CasadiLike(self.array[idx])
+        else:
+            # For other cases, delegate to the CasADi object's __getitem__
+            return CasadiLike(self.array[idx])
 
     @property
     def T(self) -> "CasadiLike":
@@ -129,6 +184,68 @@ def array(*x) -> "CasadiLike":
         """
         return CasadiLike(cs.SX(*x))
 
+    @staticmethod
+    def zeros_like(x) -> CasadiLike:
+        """
+        Args:
+            x (npt.ArrayLike): matrix
+
+        Returns:
+            npt.ArrayLike: zero matrix of dimension x
+        """
+
+        kind = (
+            cs.DM
+            if (isinstance(x, CasadiLike) and isinstance(x.array, cs.DM))
+            or isinstance(x, cs.DM)
+            else cs.SX
+        )
+
+        return (
+            CasadiLike(kind.zeros(x.array.shape))
+            if isinstance(x, CasadiLike)
+            else (
+                CasadiLike(kind.zeros(x.shape))
+                if isinstance(x, (cs.SX, cs.DM))
+                else (
+                    TypeError(f"Unsupported type for zeros_like: {type(x)}")
+                    if isinstance(x, CasadiLike)
+                    else CasadiLike(kind.zeros(x.shape))
+                )
+            )
+        )
+
+    @staticmethod
+    def ones_like(x) -> CasadiLike:
+        """
+        Args:
+            x (npt.ArrayLike): matrix
+
+        Returns:
+            npt.ArrayLike: Identity matrix of dimension x
+        """
+
+        kind = (
+            cs.DM
+            if (isinstance(x, CasadiLike) and isinstance(x.array, cs.DM))
+            or isinstance(x, cs.DM)
+            else cs.SX
+        )
+
+        return (
+            CasadiLike(kind.ones(x.array.shape))
+            if isinstance(x, CasadiLike)
+            else (
+                CasadiLike(kind.ones(x.shape))
+                if isinstance(x, (cs.SX, cs.DM))
+                else (
+                    TypeError(f"Unsupported type for ones_like: {type(x)}")
+                    if isinstance(x, CasadiLike)
+                    else CasadiLike(kind.ones(x.shape))
+                )
+            )
+        )
+
 
 class SpatialMath(SpatialMath):
 
@@ -158,7 +275,7 @@ def sin(x: npt.ArrayLike) -> "CasadiLike":
         Returns:
             CasadiLike: the sin value of x
         """
-        return CasadiLike(cs.sin(x))
+        return CasadiLike(cs.sin(x.array) if isinstance(x, CasadiLike) else cs.sin(x))
 
     @staticmethod
     def cos(x: npt.ArrayLike) -> "CasadiLike":
@@ -169,7 +286,7 @@ def cos(x: npt.ArrayLike) -> "CasadiLike":
         Returns:
             CasadiLike: the cos value of x
         """
-        return CasadiLike(cs.cos(x))
+        return CasadiLike(cs.cos(x.array) if isinstance(x, CasadiLike) else cs.cos(x))
 
     @staticmethod
     def outer(x: npt.ArrayLike, y: npt.ArrayLike) -> "CasadiLike":
@@ -206,6 +323,32 @@ def horzcat(*x) -> "CasadiLike":
         y = [xi.array if isinstance(xi, CasadiLike) else xi for xi in x]
         return CasadiLike(cs.horzcat(*y))
 
+    @staticmethod
+    def stack(x: Tuple[Union[CasadiLike, npt.ArrayLike]], axis: int = 0) -> CasadiLike:
+        """
+        Args:
+            x (Tuple[Union[CasadiLike, npt.ArrayLike]]): tuple of arrays
+            axis (int): axis to stack
+
+        Returns:
+            CasadiLike: stacked array
+
+        Notes:
+            This function is here for compatibility with the numpy_like implementation.
+        """
+
+        # check that the elements size are the same
+        for i in range(0, len(x)):
+            if len(x[i].shape) == 2:
+                raise ValueError(
+                    f"All input arrays must shape[1] != 2, {x[i].shape} found"
+                )
+
+        if axis != -1 and axis != 1:
+            raise ValueError(f"Axis must be 1 or -1, {axis} found")
+
+        return SpatialMath.vertcat(*x)
+
 
 if __name__ == "__main__":
     math = SpatialMath()

src/adam/core/rbd_algorithms.py

@@ -57,7 +57,7 @@ def crba(
             if link_i.name == self.root_link:
                 # The first "real" link. The joint is universal.
                 X_p[i] = self.math.spatial_transform(
-                    self.math.factory.eye(3), self.math.factory.zeros(3, 1)
+                    self.math.factory.eye(3), self.math.factory.zeros(3)
                 )
                 Phi[i] = self.math.factory.eye(6)
             else:
@@ -334,12 +334,13 @@ def CoM_position(
         Returns:
             com (T): The CoM position
         """
-        com_pos = self.math.factory.zeros(3, 1)
+        com_pos = self.math.factory.zeros(3)
         for item in self.model.tree:
             link = item.link
             I_H_l = self.forward_kinematics(link.name, base_transform, joint_positions)
             H_link = link.homogeneous()
             # Adding the link transform
+
             I_H_l = I_H_l @ H_link
             com_pos += I_H_l[:3, 3] * link.inertial.mass
         com_pos /= self.get_total_mass()
@@ -408,6 +409,7 @@ def rnea(
         Returns:
             tau (T): generalized force variables
         """
+
         # TODO: add accelerations
         tau = self.math.factory.zeros(self.NDoF + 6, 1)
         model_len = self.model.N
@@ -453,7 +455,7 @@ def rnea(
             if link_i.name == self.root_link:
                 # The first "real" link. The joint is universal.
                 X_p[i] = self.math.spatial_transform(
-                    self.math.factory.eye(3), self.math.factory.zeros(3, 1)
+                    self.math.factory.eye(3), self.math.factory.zeros(3)
                 )
                 Phi[i] = self.math.factory.eye(6)
                 v[i] = B_X_BI @ base_velocity