restructured integrators, added midpoint (#16)

* restructured integrators, added midpoint

* modified rollout example

* modified the integrators, added doc strings

* modified the integrators, added doc strings

* removed example 01

* Delete examples/01_intro.ipynb

* restorted example 01 from master

* added example 01

.gitignore

@@ -159,4 +159,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-*TODO
+*TODO
+.vscode*

darli/backend/base.py

@@ -263,8 +263,8 @@ def cone(
     @abstractmethod
     def integrate_configuration(
         self,
-        dt: float,
         q: ArrayLike | None = None,
         v: ArrayLike | None = None,
+        dt: float = 1,
     ) -> ArrayLike:
         ...

darli/backend/casadi.py

@@ -374,9 +374,9 @@ def cone(
 
     def integrate_configuration(
         self,
-        dt: float | cs.SX,
         q: ArrayLike | None = None,
         v: ArrayLike | None = None,
+        dt: float | cs.SX = 1.0,
     ) -> ArrayLike:
         return self.__kindyn.integrate()(
             q=q if q is not None else self._q,

darli/backend/pinocchio.py

@@ -485,9 +485,9 @@ def cone(
 
     def integrate_configuration(
         self,
-        dt: float,
         q: ArrayLike | None = None,
         v: ArrayLike | None = None,
+        dt: float = 1.0,
     ) -> ArrayLike:
         return pin.integrate(
             self.__model,

darli/modeling/base.py

@@ -359,6 +359,10 @@ def update(
 
 
 class StateSpaceBase(ABC):
+    @property
+    def integrator(self):
+        pass
+
     @property
     @abstractmethod
     def model(self) -> ModelBase:

darli/modeling/functional/state_space.py

@@ -1,7 +1,8 @@
 from ..state_space import CasadiStateSpace
 from ..base import ModelBase, StateSpaceBase, ArrayLike
 import casadi as cs
-from ..integrators import Integrator, ForwardEuler
+from ..integrators import Integrator
+from typing import Type
 
 
 class FunctionalStateSpace(StateSpaceBase):
@@ -15,6 +16,18 @@ def __init__(
         else:
             self.__space = CasadiStateSpace(model)
 
+    def set_integrator(self, integrator_cls: Type[Integrator]):
+        """
+        Set the integrator for the system using the provided Integrator class.
+
+        Args:
+            integrator_cls (Type[Integrator]): The class (constructor) of the integrator to be used.
+
+        Returns:
+            The result of setting the integrator in the underlying state space.
+        """
+        return self.__space.set_integrator(integrator_cls)
+
     @classmethod
     def from_space(cls, space: CasadiStateSpace):
         return cls(space=space)
@@ -88,11 +101,7 @@ def time_variation(self):
         )
 
     def rollout(
-        self,
-        dt: float,
-        n_steps: int,
-        control_sampling: float | None = None,
-        integrator: Integrator = ForwardEuler,
+        self, dt: float, n_steps: int, control_sampling: float | None = None
     ) -> cs.Function:
         if control_sampling is None:
             control_sampling = dt
@@ -117,12 +126,7 @@ def rollout(
             ],
             [
                 self.__space.rollout(
-                    self.__space.state,
-                    container,
-                    dt,
-                    n_steps,
-                    control_sampling,
-                    integrator,
+                    self.__space.state, container, dt, n_steps, control_sampling
                 )
             ],
             [

darli/modeling/integrators/__init__.py

@@ -2,3 +2,4 @@
 from .base import Integrator
 from .fwd_euler import ForwardEuler
 from .rk4 import RK4
+from .mid_point import MidPoint

darli/modeling/integrators/base.py

@@ -1,17 +1,101 @@
-"""Module for base class of integrators"""
+"""Module for base class of integrators."""
 
-from abc import ABC, abstractstaticmethod
-from ..base import StateSpaceBase
+from abc import ABC, abstractmethod
+from ..base import ModelBase
 from ...arrays import ArrayLike
 import casadi as cs
+from typing import Callable
 
 
-class Integrator(ABC):
-    @abstractstaticmethod
-    def forward(
-        state_space: StateSpaceBase,
-        x0: ArrayLike,
-        qfrc_u: ArrayLike,
-        dt: float | cs.SX,
-    ) -> ArrayLike:
-        pass
+class IntegratorBase(ABC):
+    """
+    Abstract base class for integrators.
+    This class provides the interface that all derived integrator classes should implement.
+    """
+
+    def __init__(self, model: ModelBase):
+        """
+        Initialize the IntegratorBase with a model of the dynamic system.
+
+        Args:
+            model (ModelBase): The model on which the integration will be performed.
+        """
+        # simplify notation for integration on manifold
+
+    @abstractmethod
+    def tangent_step(self, x: ArrayLike, tang_x: ArrayLike) -> ArrayLike:
+        """
+        Abstract method to perform step assuming the configuration is on a manifold.
+
+        Args:
+            x (ArrayLike): The current state vector of the system.
+            tang_x (ArrayLike): The tangent vector at state x.
+
+        Returns:
+            ArrayLike: The state vector after performing the tangent step.
+        """
+
+    @abstractmethod
+    def forward(self, x0: ArrayLike, u: ArrayLike, dt: float | cs.SX) -> ArrayLike:
+        """
+        Abstract method to perform a forward integration step.
+
+        Args:
+            derivative (Callable[[ArrayLike, ArrayLike, ArrayLike], ArrayLike]):
+                A function that computes the tangent of the state.
+            x0 (ArrayLike):
+                The initial state from which to start the integration.
+            u (ArrayLike): The input forces affecting the system's dynamics.
+            dt (float | cs.SX):
+                The timestep over which to integrate.
+
+        Returns:
+            ArrayLike: The state vector after performing the forward integration.
+        """
+
+    def derivative(self, x: ArrayLike, u: ArrayLike) -> ArrayLike:
+        """Calculate the derivative (tangent) of the state vector for particular model"""
+
+
+class Integrator(IntegratorBase):
+    def __init__(self, model: ModelBase):
+        """
+        Initialize the Integrator with a dynamic system model.
+
+        Args:
+            model (ModelBase): The model on which the integration will be performed.
+        """
+        super().__init__(model)
+        self.__model = model
+        self.nq = self.__model.nq
+        self.nv = self.__model.nv
+        self.__manifold_integrate = self.__model.backend.integrate_configuration
+
+    def tangent_step(self, x: ArrayLike, tang_x: ArrayLike) -> ArrayLike:
+        """
+        Concrete implementation of the tangent step for manifold integration.
+
+        Args:
+            x (ArrayLike): The current state vector of the system.
+            tang_x (ArrayLike): The tangent vector at state x.
+
+        Returns:
+            ArrayLike: The state vector after performing the tangent step.
+        """
+        # Create a container for the new state
+        container = self.__model.backend.math.zeros(self.nq + self.nv)
+        # Integrate configuration on the manifold
+        container[: self.nq] = self.__manifold_integrate(
+            q=x[: self.nq], v=tang_x[: self.nv]
+        )
+        # Velocity and acceleration are simply updated
+        container[self.nq :] = x[self.nq :] + tang_x[self.nv :]
+
+        return container.array
+
+    def derivative(self, x: ArrayLike, u: ArrayLike) -> ArrayLike:
+        q, v = x[: self.nq], x[self.nq :]
+        container = self.__model.backend.math.zeros(2 * self.nv)
+        container[: self.nv] = v
+        container[self.nv :] = self.__model.forward_dynamics(q, v, u)
+        return container.array

darli/modeling/integrators/fwd_euler.py

@@ -1,31 +1,53 @@
-"""Module for forward euler integrator"""
-from .base import Integrator, StateSpaceBase, ArrayLike, cs
+"""Module for forward euler integrator."""
+
+from .base import Integrator, ModelBase, ArrayLike, cs
+from typing import Callable
 
 
 class ForwardEuler(Integrator):
-    @staticmethod
+    """
+    Implements the Euler method for numerical integration.
+
+    The Euler method is a simple first-order integration method that updates
+    the state by taking a single step forward using the derivative of the
+    state at the current position. This method can be used for systems
+    evolving on manifolds and is computationally less intensive than higher-order methods.
+    """
+
+    def __init__(self, model: ModelBase):
+        """
+        Initialize the Euler integrator with a model that defines system dynamics,
+        which may include evolution on a manifold.
+
+        Args:
+            model (ModelBase): The DARLI model providing the system dynamics.
+        """
+        super().__init__(model)
+
     def forward(
-        state_space: StateSpaceBase,
+        self,
         x0: ArrayLike,
-        qfrc_u: ArrayLike,
+        u: ArrayLike,
         dt: cs.SX | float,
-    ):
-        nq = state_space.model.nq
-        nv = state_space.model.nv
+    ) -> ArrayLike:
+        """
+        Perform a single Euler integration step.
 
-        q, v = x0[:nq], x0[nq:]
-        vdot = state_space.model.forward_dynamics(q, v, qfrc_u)
+        Args:
+            derivative: A function that computes the tangent of the state.
+            x0: Initial state vector, which may include manifold-valued components.
+            u: Inputs forces acting on the system.
+            dt: Time step for integration.
 
-        # forward euler to integrate velocity
-        integrated_v = v + dt * vdot
+        Returns:
+            The estimated state vector after the Euler integration step.
+        """
 
-        # pinocchio fancy lie-group integration
-        integrated_q = state_space.model.backend.integrate_configuration(
-            dt, q, integrated_v
-        )
+        # Calculate the derivative at the current position
+        log = self.derivative(x0, u)
 
-        container = state_space.model.backend.math.zeros(nq + nv).array
-        container[:nq] = integrated_q
-        container[nq:] = integrated_v
+        # Update the state on manifold by taking a step along the tangent
+        # and projecting back onto the manifold
+        exp = self.tangent_step(x0, dt * log)
 
-        return container
+        return exp

darli/modeling/integrators/mid_point.py

@@ -0,0 +1,54 @@
+"""Module for forward euler integrator"""
+from .base import Integrator, ModelBase, ArrayLike, cs
+from typing import Callable
+
+
+class MidPoint(Integrator):
+    """
+    Implements the Midpoint method for numerical integration.
+
+    The Midpoint method is a second-order method that computes the derivative
+    at the midpoint of the interval to update the state. It provides a better
+    approximation than the Euler method with modest computational requirements.
+    This method can handle systems with state spaces that include manifolds.
+    """
+
+    def __init__(self, model: ModelBase):
+        """
+        Initialize the Midpoint integrator with a model that defines system dynamics,
+        which may include manifold evolution.
+
+        Args:
+            model (ModelBase): The DARLI model providing the system dynamics.
+        """
+        super().__init__(model)
+
+    def forward(
+        self,
+        x0: ArrayLike,
+        u: ArrayLike,
+        dt: cs.SX | float,
+    ) -> ArrayLike:
+        """
+        Perform a single Midpoint integration step.
+
+        Args:
+            derivative: A function that computes the tangent of the state.
+            x0: Initial state vector, which may include manifold-valued components.
+            u: Inputs forces acting on the system.
+            dt: Time step for integration.
+
+        Returns:
+            The estimated state vector after the Midpoint integration step.
+        """
+        # Compute the first time derivative of the state (Euler step)
+        k1_log = self.derivative(x0, u)
+
+        # Compute the state at the mid-point using the initial slope (k1_log)
+        k2_exp = self.tangent_step(x0, 0.5 * dt * k1_log)
+
+        # Compute the time derivative at the mid-point, based on the estimated state (k2_exp)
+        k2_log = self.derivative(k2_exp, u)
+
+        # Perform the full step using the slope at the mid-point (k2_log)
+        return self.tangent_step(x0, dt * k2_log)