Add aba function

src/adam/casadi/casadi_like.py

@@ -307,6 +307,65 @@ def transpose(self, x: CasadiLike, dims: tuple) -> CasadiLike:
         return CasadiLike(x.array.T)
 
     # --- algebra shortcuts used by algorithms ---
+    @staticmethod
+    def expand_dims(x: CasadiLike, axis: int) -> CasadiLike:
+        """Expand dimensions of a CasADi array.
+
+        Args:
+            x: Input array (CasadiLike)
+            axis: Position where new axis is to be inserted
+
+        Returns:
+            CasadiLike: Array with expanded dimensions
+        """
+        arr = x.array
+        # For CasADi, we reshape to add a dimension
+        # If axis=-1, we're adding a column dimension to make it (n,1)
+        if axis == -1:
+            # Reshape to column vector
+            return CasadiLike(cs.reshape(arr, (-1, 1)))
+        else:
+            # For other axes, just return as is (CasADi is 2D only)
+            return x
+
+    @staticmethod
+    def inv(x: CasadiLike) -> CasadiLike:
+        """Matrix inversion for CasADi.
+
+        Args:
+            x: Matrix to invert (CasadiLike)
+
+        Returns:
+            CasadiLike: Inverse of x
+        """
+        return CasadiLike(cs.inv(x.array))
+
+    @staticmethod
+    def solve(A: CasadiLike, B: CasadiLike) -> CasadiLike:
+        """Solve linear system Ax = B for x using CasADi.
+
+        Args:
+            A: Coefficient matrix (CasadiLike)
+            B: Right-hand side vector or matrix (CasadiLike)
+
+        Returns:
+            CasadiLike: Solution x
+        """
+        return CasadiLike(cs.solve(A.array, B.array))
+
+    @staticmethod
+    def mtimes(A: CasadiLike, B: CasadiLike) -> CasadiLike:
+        """Matrix-matrix multiplication for CasADi.
+
+        Args:
+            A: First matrix (CasadiLike)
+            B: Second matrix (CasadiLike)
+
+        Returns:
+            CasadiLike: Result of A @ B
+        """
+        return CasadiLike(cs.mtimes(A.array, B.array))
+
     @staticmethod
     def mxv(m: CasadiLike, v: CasadiLike) -> CasadiLike:
         """Matrix-vector multiplication for CasADi.

src/adam/casadi/computations.py

@@ -467,6 +467,83 @@ def gravity_term(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
             self.g,
         ).array
 
+    def aba(
+        self,
+        base_transform: cs.SX,
+        joint_positions: cs.SX,
+        base_velocity: cs.SX,
+        joint_velocities: cs.SX,
+        joint_torques: cs.SX,
+        external_wrenches: dict[str, cs.SX] | None = None,
+    ) -> cs.SX:
+        """Featherstone Articulated-Body Algorithm (floating base, O(n)).
+
+        Args:
+            base_transform (cs.SX): The homogenous transform from base to world frame
+            joint_positions (cs.SX): The joints position
+            base_velocity (cs.SX): The base velocity
+            joint_velocities (cs.SX): The joint velocities
+            joint_torques (cs.SX): The joint torques
+            external_wrenches (dict[str, cs.SX], optional): External wrenches applied to the robot. Defaults to None.
+
+        Returns:
+            cs.SX: The joint accelerations and the base acceleration
+        """
+        if (
+            isinstance(base_transform, cs.MX)
+            and isinstance(joint_positions, cs.MX)
+            and isinstance(base_velocity, cs.MX)
+            and isinstance(joint_velocities, cs.MX)
+            and isinstance(joint_torques, cs.MX)
+        ):
+            raise ValueError(
+                "You are using casadi MX. Please use the function KinDynComputations.aba_fun()"
+            )
+
+        return self.rbdalgos.aba(
+            base_transform,
+            joint_positions,
+            base_velocity,
+            joint_velocities,
+            joint_torques,
+            self.g,
+            external_wrenches,
+        ).array
+
+    def aba_fun(self) -> cs.Function:
+        """Returns the Articulated Body Algorithm function for forward dynamics
+
+        Returns:
+            qdd (casADi function): The joint accelerations and base acceleration
+        """
+        base_transform = cs.SX.sym("H", 4, 4)
+        joint_positions = cs.SX.sym("s", self.NDoF)
+        base_velocity = cs.SX.sym("v_b", 6)
+        joint_velocities = cs.SX.sym("s_dot", self.NDoF)
+        joint_torques = cs.SX.sym("tau", self.NDoF)
+
+        qdd = self.rbdalgos.aba(
+            base_transform,
+            joint_positions,
+            base_velocity,
+            joint_velocities,
+            joint_torques,
+            self.g,
+            None,  # external_wrenches not supported in symbolic form
+        )
+        return cs.Function(
+            "qdd",
+            [
+                base_transform,
+                joint_positions,
+                base_velocity,
+                joint_velocities,
+                joint_torques,
+            ],
+            [qdd.array],
+            self.f_opts,
+        )
+
     def CoM_position(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
         """Returns the CoM position
 

src/adam/core/array_api_math.py

@@ -238,3 +238,15 @@ def expand_dims(self, x: ArrayAPILike, axis: int) -> ArrayAPILike:
     def transpose(self, x: ArrayAPILike, dims: tuple) -> ArrayAPILike:
         xp = self._xp(x.array)
         return self.factory.asarray(xp.permute_dims(x.array, dims))
+
+    def inv(self, x: ArrayAPILike) -> ArrayAPILike:
+        xp = self._xp(x.array)
+        return self.factory.asarray(xp.linalg.inv(x.array))
+
+    def mtimes(self, A: ArrayAPILike, B: ArrayAPILike) -> ArrayAPILike:
+        xp = self._xp(A.array, B.array)
+        return self.factory.asarray(xp.matmul(A.array, B.array))
+
+    def solve(self, A: ArrayAPILike, B: ArrayAPILike) -> ArrayAPILike:
+        xp = self._xp(A.array, B.array)
+        return self.factory.asarray(xp.linalg.solve(A.array, B.array))