first commit

hpi/data_collector.py

@@ -0,0 +1,131 @@
+"""Module A: ODE simulation and trajectory data collection.
+
+Implements inverted pendulum dynamics and persistent excitation signal
+generation for data-driven adaptive optimal control.
+"""
+
+import numpy as np
+from scipy.integrate import solve_ivp
+
+
+class DataCollector:
+    """Inverted pendulum simulator with data collection capabilities.
+
+    The pendulum dynamics are given by Eq. (27):
+        dx₁/dt = x₂
+        dx₂/dt = (mgl/J)·sin(x₁) + (1/J)·u
+    """
+
+    def __init__(self, J: float = 1.0, mgl: float = 1.0, n_freqs: int = 100):
+        """
+        Args:
+            J: Moment of inertia (default 1.0).
+            mgl: Mass × gravity × length product (default 1.0).
+            n_freqs: Number of frequency components in PE signal.
+        """
+        self.J = float(J)
+        self.mgl = float(mgl)
+        self.n_freqs = n_freqs
+
+        # Pre-compute PE signal parameters
+        self._rng = np.random.RandomState(42)
+        self._omegas = self._rng.uniform(0.5, 50.0, size=n_freqs)
+        self._a = self._rng.uniform(-1.0, 1.0, size=n_freqs)
+        self._b = self._rng.uniform(-1.0, 1.0, size=n_freqs)
+        # Scale amplitudes so that max |pe_signal| ≈ 0.1
+        raw_max = np.sum(np.abs(self._a) + np.abs(self._b))
+        self._scale = 0.1 / raw_max if raw_max > 0 else 1.0
+
+    def pe_signal(self, t):
+        """Persistent excitation signal.
+
+        u_e(t) = scale · Σ (a_k·sin(ω_k·t) + b_k·cos(ω_k·t))
+
+        Args:
+            t: Time value (scalar).
+
+        Returns:
+            float: PE signal value at time t.
+        """
+        t = float(t)
+        signals = self._a * np.sin(self._omegas * t) + self._b * np.cos(self._omegas * t)
+        return self._scale * np.sum(signals)
+
+    def f(self, x):
+        """Drift dynamics f(x) — Eq. (27).
+
+        f(x) = [x₂, (mgl/J)·sin(x₁)]ᵀ
+
+        Args:
+            x: State vector (2,) or (N, 2).
+
+        Returns:
+            np.ndarray: Drift dynamics values.
+        """
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim == 1:
+            return np.array([x[1], (self.mgl / self.J) * np.sin(x[0])])
+        return np.column_stack([x[:, 1], (self.mgl / self.J) * np.sin(x[:, 0])])
+
+    def g(self, x):
+        """Input matrix g(x) — Eq. (27).
+
+        g(x) = [0, 1/J]ᵀ  (constant)
+
+        Args:
+            x: State vector (2,) or (N, 2) — ignored, kept for interface consistency.
+
+        Returns:
+            np.ndarray: Input matrix [0, 1/J]ᵀ.
+        """
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim == 1:
+            return np.array([0.0, 1.0 / self.J])
+        N = x.shape[0]
+        result = np.zeros((N, 2))
+        result[:, 1] = 1.0 / self.J
+        return result
+
+    def dynamics(self, t, x, u_func):
+        """ODE right-hand side: ẋ = f(x) + g(x)·u(t).
+
+        Args:
+            t: Current time.
+            x: State vector (2,).
+            u_func: Callable u_func(t, x) → control input.
+
+        Returns:
+            np.ndarray: State derivative (2,).
+        """
+        u_val = u_func(t, x)
+        return self.f(x) + self.g(x) * u_val
+
+    def collect_trajectory(self, x0, T, dt, u_func):
+        """Collect a trajectory by simulating the system.
+
+        Args:
+            x0: Initial state (2,).
+            T: Simulation duration.
+            dt: Time step for output sampling.
+            u_func: Control function u_func(t, x) → float.
+
+        Returns:
+            tuple: (t, X, U) where
+                t: Time points (M,).
+                X: State trajectory (M, 2).
+                U: Control inputs (M,).
+        """
+        t_eval = np.arange(0, T, dt)
+        t_span = (0.0, T)
+
+        def ode_rhs(t, x):
+            return self.dynamics(t, x, u_func)
+
+        sol = solve_ivp(ode_rhs, t_span, x0, method="RK45",
+                        t_eval=t_eval, rtol=1e-8, atol=1e-10)
+
+        t = sol.t
+        X = sol.y.T  # (M, 2)
+        U = np.array([u_func(ti, xi) for ti, xi in zip(t, X)])
+
+        return t, X, U