108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
"""Tests for DataCollector (Module A)."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from hpi.data_collector import DataCollector
|
|
|
|
|
|
@pytest.fixture
|
|
def dc():
|
|
return DataCollector(J=1.0, mgl=1.0)
|
|
|
|
|
|
class TestDataCollector:
|
|
"""Test suite for DataCollector."""
|
|
|
|
def test_f_zero(self, dc):
|
|
"""f(0) = 0 at equilibrium."""
|
|
f0 = dc.f(np.zeros(2))
|
|
np.testing.assert_allclose(f0, [0.0, 0.0], atol=1e-12)
|
|
|
|
def test_f_pi_half(self, dc):
|
|
"""f([π/2, 0]) = [0, mgl/J]."""
|
|
x = np.array([np.pi / 2, 0.0])
|
|
f_val = dc.f(x)
|
|
expected = np.array([0.0, dc.mgl / dc.J])
|
|
np.testing.assert_allclose(f_val, expected, atol=1e-12)
|
|
|
|
def test_f_batch(self, dc):
|
|
"""Batch f output shape (N, 2)."""
|
|
x = np.random.randn(10, 2)
|
|
f_vals = dc.f(x)
|
|
assert f_vals.shape == (10, 2)
|
|
|
|
def test_g_constant(self, dc):
|
|
"""g(x) = [0, 1/J] constant."""
|
|
g_val = dc.g(np.array([0.5, 0.3]))
|
|
expected = np.array([0.0, 1.0 / dc.J])
|
|
np.testing.assert_allclose(g_val, expected, atol=1e-12)
|
|
|
|
def test_g_batch(self, dc):
|
|
"""Batch g output shape (N, 2)."""
|
|
x = np.random.randn(10, 2)
|
|
g_vals = dc.g(x)
|
|
assert g_vals.shape == (10, 2)
|
|
# All rows should be [0, 1/J]
|
|
np.testing.assert_allclose(g_vals[:, 1], np.ones(10) / dc.J, atol=1e-12)
|
|
|
|
def test_pe_signal_bounded(self, dc):
|
|
"""PE signal amplitude should be <= 0.1."""
|
|
for t in np.linspace(0, 10, 100):
|
|
val = dc.pe_signal(t)
|
|
assert abs(val) <= 0.1 + 1e-12
|
|
|
|
def test_pe_signal_nonzero(self, dc):
|
|
"""PE signal should not be identically zero."""
|
|
vals = [dc.pe_signal(t) for t in np.linspace(0, 10, 100)]
|
|
assert np.max(np.abs(vals)) > 1e-6
|
|
|
|
def test_dynamics_manual(self, dc):
|
|
"""Verify ẋ = f(x) + g(x)·u."""
|
|
x = np.array([0.5, 0.3])
|
|
u = 0.7
|
|
dx = dc.dynamics(0.0, x, lambda t, x: u)
|
|
expected = dc.f(x) + dc.g(x) * u
|
|
np.testing.assert_allclose(dx, expected, atol=1e-12)
|
|
|
|
def test_collect_trajectory_shape(self, dc):
|
|
"""Trajectory outputs have consistent dimensions."""
|
|
x0 = np.array([0.5, 0.0])
|
|
T = 1.0
|
|
dt = 0.01
|
|
|
|
def u_const(t, x):
|
|
return 0.0
|
|
|
|
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
|
|
|
M_expected = int(T / dt)
|
|
assert t.shape == (M_expected,)
|
|
assert X.shape == (M_expected, 2)
|
|
assert U.shape == (M_expected,)
|
|
|
|
def test_collect_trajectory_start(self, dc):
|
|
"""First state should equal x0."""
|
|
x0 = np.array([0.5, 0.0])
|
|
T = 1.0
|
|
dt = 0.01
|
|
|
|
def u_const(t, x):
|
|
return 0.0
|
|
|
|
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
|
np.testing.assert_allclose(X[0], x0, atol=1e-10)
|
|
|
|
def test_collect_trajectory_drift(self, dc):
|
|
"""With zero control, state should evolve (pendulum swings)."""
|
|
x0 = np.array([0.5, 0.0])
|
|
T = 2.0
|
|
dt = 0.01
|
|
|
|
def u_const(t, x):
|
|
return 0.0
|
|
|
|
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
|
# Should not stay at x0 (pendulum moves)
|
|
assert np.linalg.norm(X[-1] - x0) > 1e-6
|