Add tests for wpimath.system

Author: virtuald

subprojects/robotpy-wpimath/tests/test_system.py

@@ -1,6 +1,188 @@
-import wpimath.system
-import wpimath.system.plant
+import importlib.util
+import math
 
+import pytest
 
-def test_todo():
-    pass
+import wpimath.system as system
+
+
+if importlib.util.find_spec("numpy") is None:
+    pytest.skip("numpy is not available", allow_module_level=True)
+
+import numpy as np
+
+
+def test_rk4_exponential():
+    y0 = np.array([[0.0]])
+
+    y1 = system.RK4(lambda x: np.array([[math.exp(x[0, 0])]]), y0, 0.1)
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rk4_exponential_with_u():
+    y0 = np.array([[0.0]])
+
+    y1 = system.RK4(
+        lambda x, u: np.array([[math.exp(u[0, 0] * x[0, 0])]]),
+        y0,
+        np.array([[1.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rk4_time_varying():
+    y0 = np.array([[12.0 * math.exp(5.0) / math.pow(math.exp(5.0) + 1.0, 2.0)]])
+
+    y1 = system.RK4(
+        lambda t, x: np.array([[x[0, 0] * (2.0 / (math.exp(t) + 1.0) - 1.0)]]),
+        5.0,
+        y0,
+        1.0,
+    )
+
+    expected = 12.0 * math.exp(6.0) / math.pow(math.exp(6.0) + 1.0, 2.0)
+    assert math.isclose(y1[0, 0], expected, abs_tol=1e-3)
+
+
+def test_rkdp_zero():
+    y1 = system.RKDP(
+        lambda x, u: np.zeros((1, 1)),
+        np.array([[0.0]]),
+        np.array([[0.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], 0.0, abs_tol=1e-3)
+
+
+def test_rkdp_exponential():
+    y0 = np.array([[0.0]])
+
+    y1 = system.RKDP(
+        lambda x, u: np.array([[math.exp(x[0, 0])]]),
+        y0,
+        np.array([[0.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rkdp_time_varying():
+    y0 = np.array([[12.0 * math.exp(5.0) / math.pow(math.exp(5.0) + 1.0, 2.0)]])
+
+    y1 = system.RKDP(
+        lambda t, x: np.array([[x[0, 0] * (2.0 / (math.exp(t) + 1.0) - 1.0)]]),
+        5.0,
+        y0,
+        1.0,
+        1e-12,
+    )
+
+    expected = 12.0 * math.exp(6.0) / math.pow(math.exp(6.0) + 1.0, 2.0)
+    assert math.isclose(y1[0, 0], expected, abs_tol=1e-3)
+
+
+def test_numerical_jacobian():
+    a = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+            [1.0, 1.0, 3.0, 7.0],
+        ]
+    )
+
+    def ax_fn(x):
+        return a @ x
+
+    new_a = system.numericalJacobian(ax_fn, np.zeros((4, 1)))
+    np.testing.assert_allclose(new_a, a, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_u_square():
+    a = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+            [1.0, 1.0, 3.0, 7.0],
+        ]
+    )
+    b = np.array([[1.0, 1.0], [2.0, 1.0], [3.0, 2.0], [3.0, 7.0]])
+
+    def axbu_fn(x, u, *args):
+        return a @ x + b @ u
+
+    x0 = np.zeros((4, 1))
+    u0 = np.zeros((2, 1))
+    new_a = system.numericalJacobianX(axbu_fn, x0, u0)
+    new_b = system.numericalJacobianU(axbu_fn, x0, u0)
+    np.testing.assert_allclose(new_a, a, rtol=1e-6, atol=1e-5)
+    np.testing.assert_allclose(new_b, b, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_u_rectangular():
+    c = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+        ]
+    )
+    d = np.array([[1.0, 1.0], [2.0, 1.0], [3.0, 2.0]])
+
+    def cxdu_fn(x, u, *args):
+        return c @ x + d @ u
+
+    x0 = np.zeros((4, 1))
+    u0 = np.zeros((2, 1))
+    new_c = system.numericalJacobianX(cxdu_fn, x0, u0)
+    new_d = system.numericalJacobianU(cxdu_fn, x0, u0)
+    np.testing.assert_allclose(new_c, c, rtol=1e-6, atol=1e-5)
+    np.testing.assert_allclose(new_d, d, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_passes_extra_args():
+    a = np.array([[2.0, -1.0], [0.5, 3.0]])
+    b = np.array([[1.0], [4.0]])
+    x0 = np.zeros((2, 1))
+    u0 = np.zeros((1, 1))
+
+    seen = {}
+
+    def axbu_fn(x, u, args):
+        seen["args"] = args
+        scale, bias = args
+        return scale * (a @ x) + bias * (b @ u)
+
+    new_a = system.numericalJacobianX(axbu_fn, x0, u0, 2.5, -3.0)
+
+    assert "args" in seen
+    normalized = seen["args"][0] if len(seen["args"]) == 1 else seen["args"]
+    assert normalized == (2.5, -3.0)
+    np.testing.assert_allclose(new_a, 2.5 * a, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_u_passes_extra_args():
+    a = np.array([[1.0, 0.0], [0.0, -2.0]])
+    b = np.array([[1.5], [-0.5]])
+    x0 = np.zeros((2, 1))
+    u0 = np.zeros((1, 1))
+
+    seen = {}
+
+    def axbu_fn(x, u, args):
+        seen["args"] = args
+        scale, bias = args
+        return scale * (a @ x) + bias * (b @ u)
+
+    new_b = system.numericalJacobianU(axbu_fn, x0, u0, 4.0, 0.25)
+
+    assert "args" in seen
+    normalized = seen["args"][0] if len(seen["args"]) == 1 else seen["args"]
+    assert normalized == (4.0, 0.25)
+    np.testing.assert_allclose(new_b, 0.25 * b, rtol=1e-6, atol=1e-5)