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
|