first commit

tests/test_data_collector.py

@@ -0,0 +1,107 @@
+"""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