129 lines
4.5 KiB
Python
129 lines
4.5 KiB
Python
"""Tests for NNTrainer (replaces test_solver.py)."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from hpi.nn_trainer import NNTrainer
|
|
|
|
|
|
@pytest.fixture
|
|
def trainer():
|
|
return NNTrainer(R=1.0, lr=1e-3, epochs=50, lambda_weight=1e-6)
|
|
|
|
|
|
class TestNNTrainer:
|
|
"""Test suite for NNTrainer."""
|
|
|
|
# ── integrate_interval ──
|
|
|
|
def test_integrate_interval_scalar(self):
|
|
"""Trapezoidal integration: [0,1,2] with dt=1 -> [0.5, 1.5]."""
|
|
f = np.array([0.0, 1.0, 2.0])
|
|
result = NNTrainer.integrate_interval(f, dt=1.0)
|
|
expected = np.array([0.5, 1.5])
|
|
np.testing.assert_allclose(result, expected, atol=1e-12)
|
|
|
|
def test_integrate_interval_vector(self):
|
|
"""Vector trapezoidal integration."""
|
|
f = np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]])
|
|
result = NNTrainer.integrate_interval(f, dt=1.0)
|
|
expected = np.array([[1.0, 2.0], [3.0, 4.0]])
|
|
np.testing.assert_allclose(result, expected, atol=1e-12)
|
|
|
|
# ── compute_u_hat / compute_v_hat ──
|
|
|
|
def test_compute_u_hat_manual(self, trainer):
|
|
"""u_hat returns correct shape for single and batch input."""
|
|
x_single = np.array([0.5, 0.3])
|
|
u = trainer.compute_u_hat(x_single)
|
|
assert np.isscalar(u)
|
|
assert np.isfinite(u)
|
|
|
|
x_batch = np.random.randn(10, 2).astype(np.float32)
|
|
u_batch = trainer.compute_u_hat(x_batch)
|
|
assert u_batch.shape == (10,)
|
|
assert np.all(np.isfinite(u_batch))
|
|
|
|
def test_compute_v_hat(self, trainer):
|
|
"""v_hat = u - u_hat."""
|
|
u = np.array([0.5, 0.3, 0.1])
|
|
x = np.array([[0.5, 0.3], [0.1, -0.2], [0.0, 0.0]])
|
|
u_hat = trainer.compute_u_hat(x)
|
|
v_hat = trainer.compute_v_hat(u, x)
|
|
np.testing.assert_allclose(v_hat, u - u_hat, atol=1e-10)
|
|
|
|
# ── Training output shapes ──
|
|
|
|
def test_compute_I_phi_dim(self, trainer):
|
|
"""After training, verify that V can be computed (replaces I_phi test)."""
|
|
M = 100
|
|
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
V = trainer.compute_V(x)
|
|
assert V.shape == (M,)
|
|
|
|
def test_compute_I_psi_dim(self, trainer):
|
|
"""u_hat output shape matches input (replaces I_psi test)."""
|
|
M = 100
|
|
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
u_hat = trainer.compute_u_hat(x)
|
|
assert u_hat.shape == (M,)
|
|
|
|
def test_compute_I_Q_dim(self, trainer):
|
|
"""Q(x) can be computed on trajectory (replaces I_Q test)."""
|
|
M = 100
|
|
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
Q = x[:, 0]**2 + x[:, 1]**2
|
|
assert Q.shape == (M,)
|
|
|
|
# ── System construction (adapted) ──
|
|
|
|
def test_build_system_dims(self, trainer):
|
|
"""Training produces a result dict with expected keys."""
|
|
M = 200
|
|
t = np.linspace(0, 2.0, M)
|
|
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
U = np.random.randn(M).astype(np.float32) * 0.1
|
|
|
|
result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
|
|
assert "loss_c" in result
|
|
assert "mse_c" in result
|
|
assert np.isfinite(result["loss_c"])
|
|
|
|
def test_solve_dim(self, trainer):
|
|
"""Training returns result dict with finite loss."""
|
|
M = 200
|
|
t = np.linspace(0, 2.0, M)
|
|
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
U = np.random.randn(M).astype(np.float32) * 0.1
|
|
|
|
result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
|
|
assert np.isfinite(result["loss_c"])
|
|
assert result["loss_c"] >= 0
|
|
|
|
def test_solve_no_nan(self, trainer):
|
|
"""Training produces no NaN in loss."""
|
|
M = 200
|
|
t = np.linspace(0, 2.0, M)
|
|
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
|
U = np.random.randn(M).astype(np.float32) * 0.1
|
|
|
|
result = trainer.train_one_iteration(t, X, U, L_i=0.0, prev_actor=None, dt=0.01)
|
|
assert not np.isnan(result["loss_c"])
|
|
assert not np.isinf(result["loss_c"])
|
|
|
|
def test_solve_with_trajectory_data(self, trainer):
|
|
"""Train on simulated pendulum data produces finite parameters."""
|
|
from hpi.data_collector import DataCollector
|
|
|
|
dc = DataCollector()
|
|
x0 = np.array([0.5, 0.0])
|
|
|
|
def u_func(t, x):
|
|
return dc.pe_signal(t)
|
|
|
|
t, X, U = dc.collect_trajectory(x0, T=2.0, dt=0.01, u_func=u_func)
|
|
result = trainer.train_one_iteration(t, X, U, L_i=2.0, prev_actor=None, dt=0.01)
|
|
|
|
assert np.isfinite(result["loss_c"])
|
|
assert result["loss_c"] >= 0.0
|