132 lines
3.9 KiB
Python
132 lines
3.9 KiB
Python
"""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
|